diff --git a/gateway/cache.go b/gateway/cache.go index 031d1723..a98d983d 100644 --- a/gateway/cache.go +++ b/gateway/cache.go @@ -174,9 +174,9 @@ func CacheMiddleware() gin.HandlerFunc { // Generate receipt for cache hit using current request and cached result. // Note: request_hash matches current request, response is from cache, // but both are cryptographically valid since cache key ensures identical text. - if err := generateAndSendReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, cached.Result); err != nil { + if err := streamResultAndReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, cached.Result); err != nil { log.Printf("Failed to send cached response receipt: %v", err) - // generateAndSendReceipt already sent an error response (500) + // streamResultAndReceipt already sent an error response } c.Abort() return @@ -204,14 +204,13 @@ func CacheMiddleware() gin.HandlerFunc { // Handler finished. Check status and extract result with proper locking writer.mu.RLock() statusCode := writer.ResponseWriter.Status() - bodyBytes := writer.body.Bytes() writer.mu.RUnlock() if statusCode == 200 { - // Response format: {"result": "...", "receipt": ...} - var resp map[string]interface{} - if err := json.Unmarshal(bodyBytes, &resp); err == nil { - if result, ok := resp["result"].(string); ok { + // Instead of parsing the SSE stream from bodyBytes, we retrieve the full summary + // which handleSummarize saves in the gin context. + if fullSummaryVal, exists := c.Get("full_summary"); exists { + if result, ok := fullSummaryVal.(string); ok { // Store asynchronously with a deadline to prevent indefinite goroutines go func(k, v string) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) diff --git a/gateway/cache_integration_test.go b/gateway/cache_integration_test.go index a0bf11a4..62ef828c 100644 --- a/gateway/cache_integration_test.go +++ b/gateway/cache_integration_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strconv" + "strings" "sync/atomic" "testing" "time" @@ -58,8 +59,10 @@ func TestCacheIntegration_FullFlow(t *testing.T) { aiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { aiCalls.Add(1) time.Sleep(100 * time.Millisecond) + w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(200) - w.Write([]byte(`{"choices":[{"message":{"content":"AI Summary Result"}}]}`)) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"AI Summary Result\"}}]}\n\n")) + w.Write([]byte("data: [DONE]\n\n")) })) defer aiServer.Close() @@ -198,18 +201,33 @@ func TestCacheIntegration_FullFlow(t *testing.T) { } // Verify Body - var resp1, resp2 map[string]interface{} - if err := json.Unmarshal(w1.Body.Bytes(), &resp1); err != nil { - t.Fatalf("Failed to unmarshal response 1: %v", err) - } - if err := json.Unmarshal(w2.Body.Bytes(), &resp2); err != nil { - t.Fatalf("Failed to unmarshal response 2: %v", err) + parseSSEResultForTest := func(body []byte) string { + lines := strings.Split(string(body), "\n") + var result string + for _, line := range lines { + if strings.HasPrefix(line, "data: ") { + data := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) + if data == "[DONE]" || data == "" { + continue + } + var chunk map[string]interface{} + if err := json.Unmarshal([]byte(data), &chunk); err == nil { + if text, ok := chunk["text"].(string); ok { + result += text + } + } + } + } + return result } - if val, ok := resp1["result"].(string); !ok || val != "AI Summary Result" { - t.Errorf("Unexpected result 1: %v", resp1["result"]) + result1 := parseSSEResultForTest(w1.Body.Bytes()) + result2 := parseSSEResultForTest(w2.Body.Bytes()) + + if result1 != "AI Summary Result" { + t.Errorf("Unexpected result 1: %v", result1) } - if val, ok := resp2["result"].(string); !ok || val != "AI Summary Result" { - t.Errorf("Unexpected result 2: %v", resp2["result"]) + if result2 != "AI Summary Result" { + t.Errorf("Unexpected result 2: %v", result2) } } diff --git a/gateway/errors_test.go b/gateway/errors_test.go index d1f17039..85afff5e 100644 --- a/gateway/errors_test.go +++ b/gateway/errors_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "gateway/internal/ai" "github.com/alicebob/miniredis/v2" "github.com/gin-gonic/gin" "github.com/redis/go-redis/v9" @@ -24,6 +25,10 @@ func (p failingProvider) Generate(context.Context, string) (string, error) { return "", p.err } +func (p failingProvider) GenerateStream(context.Context, string) (ai.Stream, error) { + return nil, p.err +} + func newSummarizeTestRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() diff --git a/gateway/internal/ai/ollama.go b/gateway/internal/ai/ollama.go index 57847cba..92b837fc 100644 --- a/gateway/internal/ai/ollama.go +++ b/gateway/internal/ai/ollama.go @@ -77,3 +77,58 @@ func (p *OllamaProvider) Generate(ctx context.Context, text string) (string, err return response, nil } + +type ollamaStream struct { + resp *http.Response + decoder *json.Decoder +} + +func (s *ollamaStream) Recv() (string, error) { + var chunk map[string]interface{} + if err := s.decoder.Decode(&chunk); err != nil { + return "", err // will return io.EOF when stream ends + } + if response, ok := chunk["response"].(string); ok { + return response, nil + } + return "", fmt.Errorf("invalid response from Ollama: missing response field") +} + +func (s *ollamaStream) Close() error { + return s.resp.Body.Close() +} + +func (p *OllamaProvider) GenerateStream(ctx context.Context, text string) (Stream, error) { + prompt := fmt.Sprintf("Summarize this text in 2 sentences: %s", text) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": p.model, + "prompt": prompt, + "stream": true, + }) + + req, err := http.NewRequestWithContext(ctx, "POST", p.url+"/api/generate", bytes.NewBuffer(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create Ollama request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { + return nil, context.DeadlineExceeded + } + return nil, fmt.Errorf("failed to connect to Ollama: %w", err) + } + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(body)) + } + + return &ollamaStream{ + resp: resp, + decoder: json.NewDecoder(resp.Body), + }, nil +} diff --git a/gateway/internal/ai/openrouter.go b/gateway/internal/ai/openrouter.go index 36bc48ab..be1baf4d 100644 --- a/gateway/internal/ai/openrouter.go +++ b/gateway/internal/ai/openrouter.go @@ -1,6 +1,7 @@ package ai import ( + "bufio" "bytes" "context" "encoding/json" @@ -100,3 +101,93 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string, return content, nil } + +type openRouterStream struct { + resp *http.Response + reader *bufio.Reader + sawDone bool +} + +func (s *openRouterStream) Recv() (string, error) { + for { + line, err := s.reader.ReadBytes('\n') + if err != nil { + if errors.Is(err, io.EOF) && !s.sawDone { + return "", fmt.Errorf("stream ended without [DONE] sentinel") + } + return "", err + } + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + if bytes.HasPrefix(line, []byte("data: ")) { + data := bytes.TrimPrefix(line, []byte("data: ")) + if string(data) == "[DONE]" { + s.sawDone = true + return "", io.EOF + } + var chunk map[string]interface{} + if err := json.Unmarshal(data, &chunk); err != nil { + continue // skip invalid json + } + choices, ok := chunk["choices"].([]interface{}) + if !ok || len(choices) == 0 { + continue + } + choice, ok := choices[0].(map[string]interface{}) + if !ok { + continue + } + delta, ok := choice["delta"].(map[string]interface{}) + if !ok { + continue + } + if content, ok := delta["content"].(string); ok { + return content, nil + } + } + } +} + +func (s *openRouterStream) Close() error { + return s.resp.Body.Close() +} + +func (p *OpenRouterProvider) GenerateStream(ctx context.Context, text string) (Stream, error) { + prompt := fmt.Sprintf("Summarize this text in 2 sentences: %s", text) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": p.model, + "messages": []map[string]string{ + {"role": "user", "content": prompt}, + }, + "stream": true, + }) + + req, err := http.NewRequestWithContext(ctx, "POST", p.url, bytes.NewBuffer(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create OpenRouter request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+p.apiKey) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { + return nil, context.DeadlineExceeded + } + return nil, err + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, fmt.Errorf("openrouter returned status %d: %s", resp.StatusCode, string(body)) + } + + return &openRouterStream{ + resp: resp, + reader: bufio.NewReader(resp.Body), + }, nil +} diff --git a/gateway/internal/ai/provider.go b/gateway/internal/ai/provider.go index ec027973..cc0314d6 100644 --- a/gateway/internal/ai/provider.go +++ b/gateway/internal/ai/provider.go @@ -6,10 +6,20 @@ import ( "os" ) +// Stream defines the interface for streaming AI responses +type Stream interface { + // Recv returns the next token or an error (io.EOF when done) + Recv() (string, error) + // Close releases any resources associated with the stream + Close() error +} + // Provider defines the interface for AI service providers type Provider interface { // Generate takes a context and prompt text, returns the AI-generated response Generate(ctx context.Context, prompt string) (string, error) + // GenerateStream takes a context and prompt text, returns a Stream of generated tokens + GenerateStream(ctx context.Context, prompt string) (Stream, error) } // NewProvider creates an AI provider based on the AI_PROVIDER environment variable diff --git a/gateway/main.go b/gateway/main.go index 953c94d1..0d1674ad 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -521,7 +521,7 @@ func handleSummarize(c *gin.Context) { } // 3. Call AI Service - summary, err := aiProvider.Generate(c.Request.Context(), req.Text) + stream, err := aiProvider.GenerateStream(c.Request.Context(), req.Text) if err != nil { if errors.Is(err, context.DeadlineExceeded) || errors.Is(c.Request.Context().Err(), context.DeadlineExceeded) { respondError(c, 504, "upstream_timeout", err) @@ -530,15 +530,88 @@ func handleSummarize(c *gin.Context) { } return } + defer stream.Close() - // 4. Generate & Send Receipt - if err := generateAndSendReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, summary); err != nil { - log.Printf("Failed to generate receipt: %v", err) - // generateAndSendReceipt sends error response if it fails? - // No, it returns error, we might have already written status if we aren't careful. - // Let's implement generateAndSendReceipt to handle sending response. + wantsSSE := c.GetHeader("Accept") == "text/event-stream" + var flusher http.Flusher + if wantsSSE { + c.Writer.Header().Set("Content-Type", "text/event-stream") + c.Writer.Header().Set("Cache-Control", "no-cache") + c.Writer.Header().Set("Connection", "keep-alive") + + var ok bool + flusher, ok = c.Writer.(http.Flusher) + if !ok { + respondError(c, 500, "streaming_unsupported", fmt.Errorf("streaming unsupported")) + return + } + } + + var fullSummary string + for { + text, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + if wantsSSE { + fmt.Fprintf(c.Writer, "data: {\"error\": %q}\n\n", err.Error()) + flusher.Flush() + } else { + respondError(c, 502, "upstream_stream_error", err) + } + return + } + + fullSummary += text + if wantsSSE { + chunkBytes, _ := json.Marshal(map[string]string{"text": text}) + fmt.Fprintf(c.Writer, "data: %s\n\n", chunkBytes) + flusher.Flush() + } + } + + // 4. Generate Receipt + responseMap := map[string]interface{}{ + "result": fullSummary, + } + responseBody, _ := json.Marshal(responseMap) + + receipt, err := GenerateReceipt(*paymentCtx, verifyResp.RecoveredAddress, c.Request.URL.Path, requestBody, responseBody) + if err != nil { + if wantsSSE { + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_generation_failed\"}\n\n") + flusher.Flush() + } else { + respondError(c, 500, "receipt_generation_failed", err) + } + return + } + + if err := storeReceiptWithContext(c.Request.Context(), receipt, getReceiptTTL()); err != nil { + if wantsSSE { + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_store_failed\"}\n\n") + flusher.Flush() + } else { + respondError(c, 500, "receipt_store_failed", err) + } return } + + receiptJSON, _ := json.Marshal(receipt) + receiptBase64 := base64.StdEncoding.EncodeToString(receiptJSON) + + if wantsSSE { + fmt.Fprintf(c.Writer, "data: {\"receipt\": %q}\n\n", receiptBase64) + fmt.Fprintf(c.Writer, "data: [DONE]\n\n") + flusher.Flush() + } else { + c.Writer.Header().Set("X-402-Receipt", receiptBase64) + c.Data(200, "application/json", responseBody) + } + + // Store for cache middleware + c.Set("full_summary", fullSummary) } // verifyPayment calls the verification service. @@ -609,42 +682,56 @@ func verifyPayment(ctx context.Context, signature, nonce string, timestamp uint6 return &verifyResp, &paymentCtx, nil } -// generateAndSendReceipt handles receipt generation, storage, and sending the final JSON response. -// The receipt is sent ONLY in the X-402-Receipt header, not in the response body, -// to ensure the ResponseHash in the receipt matches the actual JSON body clients receive. -func generateAndSendReceipt(c *gin.Context, paymentCtx PaymentContext, recoveredAddr string, requestBody []byte, aiResult string) error { - // Construct the response body that will be sent to client (without receipt) +// streamResultAndReceipt is used by CacheMiddleware for cache hits to stream the response identically to a live request. +func streamResultAndReceipt(c *gin.Context, paymentCtx PaymentContext, recoveredAddr string, requestBody []byte, aiResult string) error { + wantsSSE := c.GetHeader("Accept") == "text/event-stream" + + // Generate Receipt responseMap := map[string]interface{}{ "result": aiResult, } - responseBody, err := json.Marshal(responseMap) - if err != nil { - respondError(c, 500, "response_encoding_failed", err) - return err - } - - // Generate receipt with the actual response body hash + responseBody, _ := json.Marshal(responseMap) receipt, err := GenerateReceipt(paymentCtx, recoveredAddr, c.Request.URL.Path, requestBody, responseBody) if err != nil { - respondError(c, 500, "receipt_generation_failed", err) + if wantsSSE { + c.Writer.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_generation_failed\"}\n\n") + if flusher, ok := c.Writer.(http.Flusher); ok { flusher.Flush() } + } return err } if err := storeReceiptWithContext(c.Request.Context(), receipt, getReceiptTTL()); err != nil { - respondError(c, 500, "receipt_store_failed", err) + if wantsSSE { + c.Writer.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_store_failed\"}\n\n") + if flusher, ok := c.Writer.(http.Flusher); ok { flusher.Flush() } + } return err } - receiptJSON, err := json.Marshal(receipt) - if err != nil { - respondError(c, 500, "receipt_encoding_failed", err) - return err - } + receiptJSON, _ := json.Marshal(receipt) receiptBase64 := base64.StdEncoding.EncodeToString(receiptJSON) - // Send receipt in header only (not in body) so ResponseHash matches body - c.Header("X-402-Receipt", receiptBase64) - c.JSON(200, responseMap) + if wantsSSE { + c.Writer.Header().Set("Content-Type", "text/event-stream") + c.Writer.Header().Set("Cache-Control", "no-cache") + c.Writer.Header().Set("Connection", "keep-alive") + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + return fmt.Errorf("streaming unsupported") + } + + chunkBytes, _ := json.Marshal(map[string]string{"text": aiResult}) + fmt.Fprintf(c.Writer, "data: %s\n\n", chunkBytes) + fmt.Fprintf(c.Writer, "data: {\"receipt\": %q}\n\n", receiptBase64) + fmt.Fprintf(c.Writer, "data: [DONE]\n\n") + flusher.Flush() + } else { + c.Writer.Header().Set("X-402-Receipt", receiptBase64) + c.Data(200, "application/json", responseBody) + } return nil } diff --git a/gateway/middleware.go b/gateway/middleware.go index 652bc332..eb733794 100644 --- a/gateway/middleware.go +++ b/gateway/middleware.go @@ -58,12 +58,13 @@ func CorrelationIDMiddleware() gin.HandlerFunc { // decide whether to send the real response or a timeout response without // racing with handler writes. type bufferedWriter struct { - buf *bytes.Buffer - head http.Header - status int - wrote bool - closed bool - mu sync.RWMutex + buf *bytes.Buffer + head http.Header + status int + wrote bool + closed bool + flushed bool + mu sync.RWMutex } // newBufferedWriter returns an initialized bufferedWriter used to capture @@ -137,15 +138,21 @@ func (b *bufferedWriter) Status() int { // flushTo writes buffered headers and body to the real writer. func (b *bufferedWriter) flushTo(w http.ResponseWriter) { - b.mu.RLock() - defer b.mu.RUnlock() - for k, vv := range b.head { - for _, v := range vv { - w.Header().Add(k, v) + b.mu.Lock() + defer b.mu.Unlock() + if !b.flushed { + for k, vv := range b.head { + for _, v := range vv { + w.Header().Add(k, v) + } } + w.WriteHeader(b.Status()) + b.flushed = true + } + if b.buf.Len() > 0 { + _, _ = w.Write(b.buf.Bytes()) + b.buf.Reset() } - w.WriteHeader(b.Status()) - _, _ = w.Write(b.buf.Bytes()) } // RequestTimeoutMiddleware applies a context timeout to the request and @@ -245,9 +252,8 @@ func (rws *responseWriterShim) Written() bool { return rws.b func (rws *responseWriterShim) Size() int { return rws.bw.buf.Len() } func (rws *responseWriterShim) WriteHeaderNowWithoutLock() {} -// Flush flushes the response to the client if the underlying writer -// supports http.Flusher. This is a no-op otherwise. func (rws *responseWriterShim) Flush() { + rws.bw.flushTo(rws.orig) if fl, ok := rws.orig.(http.Flusher); ok { fl.Flush() } diff --git a/gateway/receipt_store_integration_test.go b/gateway/receipt_store_integration_test.go index c7a2e400..d46d15eb 100644 --- a/gateway/receipt_store_integration_test.go +++ b/gateway/receipt_store_integration_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strconv" + "strings" "sync" "testing" "time" @@ -83,11 +84,24 @@ func TestRedisReceiptStore_PersistsAcrossGatewayRestart(t *testing.T) { t.Fatalf("create receipt status=%d body=%s", createResp.Code, createResp.Body.String()) } - receiptHeader := createResp.Header().Get("X-402-Receipt") - if receiptHeader == "" { - t.Fatal("missing X-402-Receipt header") + var receiptBase64 string + lines := strings.Split(createResp.Body.String(), "\n") + for _, line := range lines { + if strings.HasPrefix(line, "data: ") { + dataStr := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) + var chunk map[string]interface{} + if err := json.Unmarshal([]byte(dataStr), &chunk); err == nil { + if r, ok := chunk["receipt"].(string); ok { + receiptBase64 = r + } + } + } + } + + if receiptBase64 == "" { + t.Fatal("missing receipt event in SSE stream") } - receiptJSON, err := base64.StdEncoding.DecodeString(receiptHeader) + receiptJSON, err := base64.StdEncoding.DecodeString(receiptBase64) if err != nil { t.Fatalf("decode receipt header: %v", err) } diff --git a/gateway/timeout_test.go b/gateway/timeout_test.go index b0e2e03e..7581c3f0 100644 --- a/gateway/timeout_test.go +++ b/gateway/timeout_test.go @@ -63,7 +63,7 @@ func TestCallOpenRouter_RespectsContextTimeout(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() - _, err = provider.Generate(ctx, "hello") + _, err = provider.GenerateStream(ctx, "hello") if err == nil { t.Fatalf("Expected timeout error from provider, got nil") } diff --git a/tests/e2e.test.ts b/tests/e2e.test.ts index 3f5ce3b7..9b0a222e 100644 --- a/tests/e2e.test.ts +++ b/tests/e2e.test.ts @@ -68,6 +68,7 @@ describe("MicroAI Paygate E2E Flow", () => { method: "POST", headers: { "Content-Type": "application/json", + "Accept": "text/event-stream", "X-402-Signature": signature, "X-402-Nonce": paymentContext.nonce, "X-402-Timestamp": paymentContext.timestamp.toString(), @@ -86,8 +87,29 @@ describe("MicroAI Paygate E2E Flow", () => { } expect(res.status).toBe(200); - const data = await res.json() as any; - expect(data.result).toBeDefined(); + const textStr = await res.text(); + expect(textStr).toContain("data: {"); + expect(textStr).toContain("[DONE]"); + + const lines = textStr.split("\n"); + let hasText = false; + let hasReceipt = false; + let hasDone = false; + for (const line of lines) { + if (!line.startsWith("data: ")) continue; + const data = line.replace("data: ", "").trim(); + if (data === "[DONE]") hasDone = true; + else { + try { + const parsed = JSON.parse(data); + if (parsed.text) hasText = true; + if (parsed.receipt) hasReceipt = true; + } catch (e) {} + } + } + expect(hasText).toBe(true); + expect(hasReceipt).toBe(true); + expect(hasDone).toBe(true); }, 30000); it("should reject replayed signed payment context", async () => { diff --git a/web/src/hooks/use-x402.ts b/web/src/hooks/use-x402.ts index 7f5933fe..dea0eb33 100644 --- a/web/src/hooks/use-x402.ts +++ b/web/src/hooks/use-x402.ts @@ -111,7 +111,9 @@ export function useX402() { if (first.status === 200) { update({ step: "receipt" }); - const { summary, receipt } = await readSummarizeSuccess(first); + const { summary, receipt } = await readSummarizeSuccess(first, (text) => { + update({ summary: text }); + }); if (receipt) saveReceipt(receipt, text); track(AnalyticsEvent.SummaryCompleted, { ...flowProps, @@ -280,7 +282,9 @@ export function useX402() { } update({ step: "receipt" }); - const { summary, receipt } = await readSummarizeSuccess(retry); + const { summary, receipt } = await readSummarizeSuccess(retry, (text) => { + update({ summary: text }); + }); if (receipt) saveReceipt(receipt, text); track(AnalyticsEvent.SummaryCompleted, { ...flowProps, diff --git a/web/src/lib/x402-client.ts b/web/src/lib/x402-client.ts index 2a24c42a..7f0a3c55 100644 --- a/web/src/lib/x402-client.ts +++ b/web/src/lib/x402-client.ts @@ -15,7 +15,7 @@ export async function postSummarize( ): Promise { return fetch(`${getGatewayUrl()}/api/ai/summarize`, { method: "POST", - headers: { "Content-Type": "application/json", ...headers }, + headers: { "Content-Type": "application/json", "Accept": "text/event-stream", ...headers }, body: JSON.stringify({ text }), }); } @@ -68,12 +68,51 @@ export type SummarizeSuccess = { receipt: SignedReceipt | null; }; -export async function readSummarizeSuccess(res: Response): Promise { - const data = (await res.json()) as { result?: string }; - const summary = data.result ?? ""; +export async function readSummarizeSuccess( + res: Response, + onChunk?: (text: string) => void +): Promise { + if (!res.body) throw new Error("No response body"); + const reader = res.body.getReader(); + const decoder = new TextDecoder("utf-8"); + let summary = ""; + let receipt: SignedReceipt | null = null; + let buffer = ""; + let sawDone = false; - const headerVal = res.headers.get("x-402-receipt") ?? res.headers.get("X-402-Receipt"); - const receipt = headerVal ? safeDecodeReceiptHeader(headerVal) : null; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + + const lines = buffer.split("\n"); + buffer = lines.pop() ?? ""; + for (const line of lines) { + if (line.startsWith("data: ")) { + const dataStr = line.slice("data: ".length).trim(); + if (dataStr === "[DONE]") { + sawDone = true; + continue; + } + try { + const parsed = JSON.parse(dataStr); + if (parsed.error) { + throw new Error(`SSE error: ${parsed.error}`); + } + if (parsed.text) { + summary += parsed.text; + if (onChunk) onChunk(summary); + } else if (parsed.receipt) { + receipt = safeDecodeReceiptHeader(parsed.receipt); + } + } catch (e) { + if (e instanceof Error && e.message.startsWith("SSE error:")) throw e; + // ignore + } + } + } + } + if (!sawDone) throw new Error("Stream ended before [DONE]"); return { summary, receipt }; }