From de017d30822577b6073454a6eb8e2b7fcb4bb42c Mon Sep 17 00:00:00 2001
From: Manuel Odendahl <wesen@ruinwesen.com>
Date: Sun, 26 Jan 2025 19:40:57 -0500
Subject: [PATCH] :sparkles: Mess around with SSE sizes and tool result sizes

---
 .../simplify-html/examples/filter-pubmed.yaml | 40 +++++++++++++++++++
 pkg/client/sse.go                             |  6 ++-
 pkg/client/stdio.go                           |  7 +++-
 pkg/cmds/shell-tool-provider.go               |  8 +++-
 pkg/server/transports/stdio/sse_bridge.go     |  2 +-
 5 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 cmd/tools/simplify-html/examples/filter-pubmed.yaml

diff --git a/cmd/tools/simplify-html/examples/filter-pubmed.yaml b/cmd/tools/simplify-html/examples/filter-pubmed.yaml
new file mode 100644
index 0000000..f9b16a7
--- /dev/null
+++ b/cmd/tools/simplify-html/examples/filter-pubmed.yaml
@@ -0,0 +1,40 @@
+selectors:
+  - type: css
+    mode: filter
+    selector: "div[role='banner']"
+  - type: css
+    mode: filter
+    selector: ".usa-banner"
+  - type: css
+    mode: filter
+    selector: ".ncbi-header"
+  - type: css
+    mode: filter
+    selector: "[role='navigation']"
+  - type: css
+    mode: filter
+    selector: ".search-form"
+  - type: css
+    mode: filter
+    selector: ".multiple-results-actions"
+  - type: css
+    mode: filter
+    selector: ".action-panel"
+  - type: css
+    mode: filter
+    selector: ".search-options"
+  - type: css
+    mode: filter
+    selector: ".sort-dropdown"
+  - type: css
+    mode: filter
+    selector: ".display-options"
+  - type: css
+    mode: filter
+    selector: ".inner-wrap .static-filters"
+  - type: css
+    mode: filter
+    selector: "footer"
+  - type: css
+    mode: filter
+    selector: ".results-amount-container"
\ No newline at end of file
diff --git a/pkg/client/sse.go b/pkg/client/sse.go
index 6f1a926..e2fc67e 100644
--- a/pkg/client/sse.go
+++ b/pkg/client/sse.go
@@ -106,6 +106,7 @@ func (t *SSETransport) Send(ctx context.Context, request *protocol.Request) (*pr
 			return nil, fmt.Errorf("server error: %s", string(event.Data))
 		}
 
+		t.logger.Debug().Int("dataLength", len(event.Data)).Msg("FOOBAR")
 		t.logger.Debug().
 			Str("event", string(event.Event)).
 			RawJSON("data", event.Data).
@@ -142,7 +143,10 @@ func (t *SSETransport) initializeSSE(ctx context.Context) error {
 		err := t.sseClient.SubscribeWithContext(subCtx, "", func(msg *sse.Event) {
 			t.logger.Debug().
 				Str("event", string(msg.Event)).
-				RawJSON("data", msg.Data).
+				Str("data", string(msg.Data)).
+				Str("retry", string(msg.Retry)).
+				Str("id", string(msg.ID)).
+				Str("comment", string(msg.Comment)).
 				Msg("Received SSE event")
 
 			// Handle endpoint event
diff --git a/pkg/client/stdio.go b/pkg/client/stdio.go
index 1d09640..3ebb8ca 100644
--- a/pkg/client/stdio.go
+++ b/pkg/client/stdio.go
@@ -26,8 +26,13 @@ type StdioTransport struct {
 
 // NewStdioTransport creates a new stdio transport
 func NewStdioTransport(logger zerolog.Logger) *StdioTransport {
+	scanner := bufio.NewScanner(os.Stdin)
+	// Set 1MB buffer size to avoid "token too long" errors
+	buf := make([]byte, 1024*1024)
+	scanner.Buffer(buf, len(buf))
+	
 	return &StdioTransport{
-		scanner: bufio.NewScanner(os.Stdin),
+		scanner: scanner,
 		writer:  json.NewEncoder(os.Stdout),
 		logger:  logger,
 	}
diff --git a/pkg/cmds/shell-tool-provider.go b/pkg/cmds/shell-tool-provider.go
index 2c49e8a..4d884b7 100644
--- a/pkg/cmds/shell-tool-provider.go
+++ b/pkg/cmds/shell-tool-provider.go
@@ -150,7 +150,13 @@ func (p *ShellToolProvider) CallTool(ctx context.Context, name string, arguments
 		return protocol.NewErrorToolResult(protocol.NewTextContent(err.Error())), nil
 	}
 
-	result := protocol.NewToolResult(protocol.WithText(buf.String()))
+	text := buf.String()
+	l := 62 * 1024
+	if len(text) > l {
+		text = text[:l]
+	}
+
+	result := protocol.NewToolResult(protocol.WithText(text))
 
 	if p.tracingDir != "" {
 		timestamp := time.Now().Format("2006-01-02T15-04-05.000")
diff --git a/pkg/server/transports/stdio/sse_bridge.go b/pkg/server/transports/stdio/sse_bridge.go
index ad36112..56b46f9 100644
--- a/pkg/server/transports/stdio/sse_bridge.go
+++ b/pkg/server/transports/stdio/sse_bridge.go
@@ -42,7 +42,7 @@ func NewSSEBridgeServer(logger zerolog.Logger, sseURL string) *SSEBridgeServer {
 	}
 
 	// Create a new logger that writes to the tagged stderr
-	taggedLogger := logger.Output(consoleWriter)
+	taggedLogger := logger.Output(consoleWriter).With().Caller().Logger()
 
 	// Strip trailing slashes from the SSE URL
 	sseURL = strings.TrimRight(sseURL, "/")