// This file implements the shared transport driver (US-007 support): a // provider-agnostic layer that turns an HTTP request into a stream of // StreamEvents. Each provider degenerates to a stateful Decoder; the transport // owns HTTP + SSE line parsing + retry + dual watchdogs + the dual failure // model. // // The design mirrors pi's providerio: the transport never returns a runtime // failure as a Go error once streaming has begun — it rides the stream as a // terminal StreamErrorEvent. Only the earliest "cannot build the stream" case // (bad request construction) is a returned error. package provider import ( "bufio" "context" "errors" "fmt" "io" "net" "net/http" "os" "strconv" "strings" "time" "github.com/smallnest/pigo/internal/agentcore" ) // StreamEvent is the transport-level alias for AssistantMessageEvent. Decoders // produce these; the transport forwards them onto the stream. (Decision #25: // reuse AssistantMessageEvent rather than a parallel event type.) type StreamEvent = AssistantMessageEvent // Decoder is the per-provider stateful SSE payload decoder. The transport calls // Decode for every complete SSE data payload (one event's worth of bytes) and // Finish once the stream ends so the decoder can flush any buffered terminal // event. type Decoder interface { // Decode turns one SSE data payload into zero or more StreamEvents. A // returned error is treated as a runtime stream failure (terminal error // event), never a panic. Decode(payload []byte) ([]StreamEvent, error) // Finish flushes any trailing state, returning a final batch of events. Finish() ([]StreamEvent, error) } // defaultIdleTimeout is the watchdog idle window; PIGO_STREAM_IDLE_TIMEOUT // (a Go duration string, e.g. "3m") overrides it. const defaultIdleTimeout = 5 * time.Minute const ( // defaultMaxConnectRetries bounds connect-only retries when TransportConfig // leaves MaxConnectRetries at zero. defaultMaxConnectRetries = 2 // statusTooManyRequestsCF (529) is Cloudflare's "site overloaded" status, // which some upstreams also emit; treated as retryable alongside 429/503. statusTooManyRequestsCF = 529 // stallFactor slackens the content-stall watchdog relative to the idle // window (stall = idle × stallFactor) so a slow-but-progressing stream is not // killed by the stall guard. stallFactor = 1.2 // errorBodyLimit bounds how many bytes of an upstream error body are read // into the returned error message. errorBodyLimit = 4096 ) // idleTimeout resolves the configured idle watchdog window. func idleTimeout() time.Duration { if v := os.Getenv("PIGO_STREAM_IDLE_TIMEOUT"); v != "" { if d, err := time.ParseDuration(v); err == nil && d > 0 { return d } } return defaultIdleTimeout } // TransportConfig configures a single StreamRequest run. type TransportConfig struct { // Client is the HTTP client; defaults to http.DefaultClient when nil. Client *http.Client // NewRequest builds a fresh *http.Request for each connection attempt. It is // called once per connect (initial + reconnects) so retries never replay a // consumed body — the caller owns idempotent request construction. NewRequest func(ctx context.Context) (*http.Request, error) // Decoder converts SSE payloads to StreamEvents (required). Decoder Decoder // MaxConnectRetries bounds connect-only retries (default 2). MaxConnectRetries int } // StreamRequest runs cfg as a transport stream. Per the dual failure model it // returns an error only when the very first request cannot be built or the // initial connection can never be established; every runtime failure once // streaming begins rides the returned stream as a terminal StreamErrorEvent. func StreamRequest(ctx context.Context, cfg TransportConfig) (*AssistantMessageEventStream, error) { if cfg.NewRequest == nil { return nil, errors.New("transport: NewRequest is required") } if cfg.Decoder == nil { return nil, errors.New("transport: Decoder is required") } client := cfg.Client if client == nil { client = http.DefaultClient } maxRetries := cfg.MaxConnectRetries if maxRetries == 0 { maxRetries = defaultMaxConnectRetries } // Connect once up front so a "cannot even build the stream" failure surfaces // as a returned error (the only early-error case per FR-13). resp, err := connect(ctx, client, cfg.NewRequest, maxRetries) if err != nil { return nil, err } stream := NewAssistantMessageEventStream(0) go pump(ctx, stream, resp, cfg.Decoder) return stream, nil } // connect performs the initial connection with retry. It only retries when the // server explicitly signals a retryable condition (429/503/529); it never // replays a consumed stream, so retrying at connect time is always safe. func connect(ctx context.Context, client *http.Client, newReq func(context.Context) (*http.Request, error), maxRetries int) (*http.Response, error) { var lastErr error for attempt := 0; attempt <= maxRetries; attempt++ { req, err := newReq(ctx) if err != nil { return nil, fmt.Errorf("transport: build request: %w", err) } resp, err := client.Do(req) if err != nil { lastErr = classifyTransportError(err) if !isRetryableNetErr(err) || attempt == maxRetries { return nil, lastErr } if !sleepBackoff(ctx, attempt, 0) { return nil, ctx.Err() } continue } if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable || resp.StatusCode == statusTooManyRequestsCF { wait := retryAfter(resp.Header) resp.Body.Close() lastErr = fmt.Errorf("transport: upstream %d", resp.StatusCode) if attempt == maxRetries { return nil, lastErr } if !sleepBackoff(ctx, attempt, wait) { return nil, ctx.Err() } continue } if resp.StatusCode >= 400 { body, _ := io.ReadAll(io.LimitReader(resp.Body, errorBodyLimit)) resp.Body.Close() return nil, fmt.Errorf("transport: upstream %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) } return resp, nil } return nil, lastErr } // pump drives the SSE read loop with dual watchdogs, decoding payloads and // forwarding events onto the stream. All runtime failures become a terminal // error event; pump always closes the stream. func pump(ctx context.Context, stream *AssistantMessageEventStream, resp *http.Response, dec Decoder) { defer stream.Close() defer resp.Body.Close() idle := idleTimeout() // content-stall watchdog is slightly slacker than idle (idle × stallFactor) // so a slow but progressing stream is not killed by the stall guard. stall := time.Duration(float64(idle) * stallFactor) // The watchdog fires by cancelling a derived context; reads race against it. watchCtx, cancel := context.WithCancel(ctx) defer cancel() // done stops the reader goroutine so it never blocks on a send after pump // returns (watchdog / abort paths), avoiding a goroutine leak. done := make(chan struct{}) defer close(done) lines := make(chan string) readErr := make(chan error, 1) go readLines(resp.Body, lines, readErr, done) var dataBuf strings.Builder idleTimer := time.NewTimer(idle) stallTimer := time.NewTimer(stall) defer idleTimer.Stop() defer stallTimer.Stop() emit := func(events []StreamEvent) bool { for _, ev := range events { if err := stream.Emit(watchCtx, ev); err != nil { return false } } return true } fail := func(msg string, err error) { stream.Emit(context.Background(), StreamErrorEvent{ Message: agentcore.AssistantMessage{ RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: msg, }, Err: err, }) } flush := func() bool { if dataBuf.Len() == 0 { return true } payload := dataBuf.String() dataBuf.Reset() if payload == "[DONE]" { return true } events, err := dec.Decode([]byte(payload)) if err != nil { fail("decode error: "+err.Error(), err) return false } return emit(events) } for { select { case <-ctx.Done(): fail("stream aborted", ctx.Err()) return case <-idleTimer.C: fail("idle timeout: no data received", errStreamIdle) return case <-stallTimer.C: fail("content stall timeout", errStreamStall) return case err := <-readErr: if err != nil && !errors.Is(err, io.EOF) { fail("read error: "+classifyTransportError(err).Error(), err) return } // Clean EOF: flush any buffered payload, then finish the decoder. if !flush() { return } finalEvents, ferr := dec.Finish() if ferr != nil { fail("finish error: "+ferr.Error(), ferr) return } emit(finalEvents) return case line, ok := <-lines: if !ok { continue } // Any byte resets the idle watchdog; a flushed event resets stall. resetTimer(idleTimer, idle) line = strings.TrimRight(line, "\r\n") switch { case line == "": // Blank line = event boundary: flush accumulated data. if !flush() { return } resetTimer(stallTimer, stall) case strings.HasPrefix(line, ":"): // Comment / keep-alive: ignore payload, watchdog already reset. case strings.HasPrefix(line, "data:"): dataBuf.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:"))) default: // Non-data field (event:, id:, etc.) — ignored for our decoders. } } } } // readLines reads the body line by line, pushing each onto lines and the final // error (io.EOF on clean close) onto readErr. It stops promptly when done is // closed so pump can return on a watchdog/abort without leaking this goroutine. func readLines(r io.Reader, lines chan<- string, readErr chan<- error, done <-chan struct{}) { br := bufio.NewReader(r) for { line, err := br.ReadString('\n') if line != "" { select { case lines <- line: case <-done: return } } if err != nil { select { case readErr <- err: case <-done: } return } } } // resetTimer stops and re-arms t to fire after d. func resetTimer(t *time.Timer, d time.Duration) { if !t.Stop() { select { case <-t.C: default: } } t.Reset(d) } // Sentinel errors for watchdog classification. var ( errStreamIdle = errors.New("stream idle timeout") errStreamStall = errors.New("stream content stall") ) // retryAfter parses a Retry-After header (seconds or HTTP-date), returning 0 // when absent/unparseable. func retryAfter(h http.Header) time.Duration { v := h.Get("Retry-After") if v == "" { return 0 } if secs, err := strconv.Atoi(v); err == nil && secs >= 0 { return time.Duration(secs) * time.Second } if t, err := http.ParseTime(v); err == nil { if d := time.Until(t); d > 0 { return d } } return 0 } // sleepBackoff waits for the retry delay (Retry-After if given, else // exponential), honoring ctx cancellation. Returns false if ctx was cancelled. func sleepBackoff(ctx context.Context, attempt int, retryAfter time.Duration) bool { d := retryAfter if d == 0 { d = time.Duration(1<