31 lines
961 B
Go
31 lines
961 B
Go
package toolsrv
|
|
|
|
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
|
|
}
|