streaming: emit partial tool output via context callback

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 {

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, &notif) == 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.

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
 	}

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
 

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
+}