diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 35d9b12..36b2546 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -10,6 +10,7 @@ import ( "time" "ollie/pkg/backend" + "ollie/pkg/tools" ) const maxTransientRetries = 3 @@ -298,7 +299,10 @@ func run(ctx context.Context, cfg agentConfig, state state) error { var resultBlocks []backend.ContentBlock var isErr bool if cfg.Exec != nil { - out, blocks, err := cfg.Exec(ctx, tc.Name, tc.Arguments) + streamCtx := tools.WithOutputStream(ctx, func(data string) { + emit(cfg, Event{Role: "exec", Name: tc.Name, Content: data}) + }) + out, blocks, err := cfg.Exec(streamCtx, tc.Name, tc.Arguments) if err != nil { isErr = true if ctx.Err() != nil { diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go index 58e19f9..56c8874 100644 --- a/pkg/remote/remote.go +++ b/pkg/remote/remote.go @@ -303,18 +303,36 @@ func (s *Server) CallTool(ctx context.Context, tool string, args json.RawMessage return nil, fmt.Errorf("remote rpc write: %w", err) } - var resp rpcResponse - if err := s.dec.Decode(&resp); err != nil { + // Read responses, handling interleaved streaming notifications. + for { + var resp rpcResponse + if err := s.dec.Decode(&resp); err != nil { + s.mu.Unlock() + return nil, fmt.Errorf("remote rpc read: %w", err) + } + + // Notification: id == 0 and method is "output". + // (JSON-RPC notifications have no id; our struct decodes missing id as 0.) + if resp.ID == 0 && resp.Result != nil { + // Streaming output notification — emit to context callback. + var notif outputNotification + if json.Unmarshal(resp.Result, ¬if) == nil && notif.Data != "" { + tools.StreamOutput(ctx, notif.Data) + } + continue + } + + // Actual response for our request. s.mu.Unlock() - return nil, fmt.Errorf("remote rpc read: %w", err) + if resp.Error != nil { + return nil, fmt.Errorf("remote error: %s", resp.Error.Message) + } + return resp.Result, nil } - s.mu.Unlock() +} - if resp.Error != nil { - return nil, fmt.Errorf("remote error: %s", resp.Error.Message) - } - - return resp.Result, nil +type outputNotification struct { + Data string `json:"data"` } // fetchHostInfo retrieves environment details from the remote host. diff --git a/pkg/tools/execute/code.go b/pkg/tools/execute/code.go index 8789a3b..0af80d8 100644 --- a/pkg/tools/execute/code.go +++ b/pkg/tools/execute/code.go @@ -166,6 +166,7 @@ type limitedWriter struct { written int limit int truncated bool + stream func(string) // if non-nil, called with each chunk of output } func (lw *limitedWriter) Write(p []byte) (n int, err error) { @@ -185,6 +186,9 @@ func (lw *limitedWriter) Write(p []byte) (n int, err error) { written, err := lw.w.Write(toWrite) lw.written += written + if lw.stream != nil { + lw.stream(string(toWrite[:written])) + } if err != nil { return written, err } diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go index ecb8c6a..2581d9b 100644 --- a/pkg/tools/execute/server.go +++ b/pkg/tools/execute/server.go @@ -466,7 +466,11 @@ func (e *Server) executeElevated(ctx context.Context, cmd, dir string, timeout i c.WaitDelay = time.Second var outputBuf bytes.Buffer - lw := &limitedWriter{w: &outputBuf, limit: 10 * 1024 * 1024} + lw := &limitedWriter{ + w: &outputBuf, + limit: 10 * 1024 * 1024, + stream: tools.StreamFunc(ctx), + } c.Stdout = lw c.Stderr = lw @@ -752,7 +756,11 @@ func (e *Server) executeWithStdin(ctx context.Context, code, language string, ti cmd.WaitDelay = time.Second var outputBuf bytes.Buffer - lw := &limitedWriter{w: &outputBuf, limit: 10 * 1024 * 1024} + lw := &limitedWriter{ + w: &outputBuf, + limit: 10 * 1024 * 1024, + stream: tools.StreamFunc(ctx), + } cmd.Stdout = lw cmd.Stderr = lw diff --git a/pkg/tools/stream.go b/pkg/tools/stream.go new file mode 100644 index 0000000..7c75f64 --- /dev/null +++ b/pkg/tools/stream.go @@ -0,0 +1,30 @@ +package tools + +import "context" + +// OutputFunc is a callback for streaming partial tool output. +type OutputFunc func(data string) + +type streamKey struct{} + +// WithOutputStream returns a context carrying an output streaming callback. +// Tool servers can call StreamOutput(ctx, data) to emit partial results. +func WithOutputStream(ctx context.Context, fn OutputFunc) context.Context { + return context.WithValue(ctx, streamKey{}, fn) +} + +// StreamOutput sends partial output to the streaming callback in ctx, if any. +func StreamOutput(ctx context.Context, data string) { + if fn, ok := ctx.Value(streamKey{}).(OutputFunc); ok && fn != nil { + fn(data) + } +} + +// StreamFunc returns the streaming callback from ctx, or nil if none set. +// Useful for passing to writers that call the function directly. +func StreamFunc(ctx context.Context) func(string) { + if fn, ok := ctx.Value(streamKey{}).(OutputFunc); ok { + return func(data string) { fn(data) } + } + return nil +}