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
}