Skip to content

Conversation

@tkattkat
Copy link
Collaborator

@tkattkat tkattkat commented Dec 28, 2025

Why

The middleware approach kept the full message data in state and created modified copies each time before passing to the LLM so compression never actually reduced messages memory.

What Changed

  • Moved processMessages from middleware to prepareStep callback, allowing us to mutate messages in place so image/aria tree compression actually takes effect in memory
  • updated compression to handle both output.value and result fields (base64 data lives in output.value)
  • Removed dead code from messageProcessing.ts

Test Plan

  • Verified data is properly removed from returned messages
  • Verified messages can still be passed from one agent run to the next with no issues

Summary by cubic

Move agent message compression to prepareStep to mutate messages in place, pruning old screenshots and ARIA trees to reduce memory. Also support both output.value and result tool payloads.

  • Refactors
    • Process messages in prepareStep; remove transformParams middleware.
    • In-place compression: keep 2 latest screenshots and 1 latest ARIA tree; replace older data with placeholders.
    • Support tool results in output.value (base64) and result fields.
    • Removed dead code from messageProcessing.ts.

Written for commit 7dd4a35. Summary will update automatically on new commits.

@changeset-bot
Copy link

changeset-bot bot commented Dec 28, 2025

🦋 Changeset detected

Latest commit: 7dd4a35

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 3 packages
Name Type
@browserbasehq/stagehand Patch
@browserbasehq/stagehand-evals Patch
@browserbasehq/stagehand-server Patch

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

Copy link
Contributor

@cubic-dev-ai cubic-dev-ai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No issues found across 3 files

@greptile-apps
Copy link
Contributor

greptile-apps bot commented Dec 28, 2025

Greptile Summary

Moved message compression from middleware to prepareStep callback to enable in-place mutation, fixing a critical memory leak where base64 screenshot/ariaTree data was never actually removed from memory.

Key changes:

  • messageProcessing.ts: Refactored processMessages() to mutate messages in-place instead of returning copies; updated to handle both output.value and result fields where base64 data lives; removed unused CompressionStats interface and calculation logic
  • v3AgentHandler.ts: Created createPrepareStep() wrapper that calls processMessages() on the messages array before each LLM step, replacing the middleware transformParams approach that created copies; added TODO comment noting middleware may no longer be needed

Impact: The middleware approach was creating modified copies before passing to the LLM, but the original state messages remained unchanged in memory. This refactor mutates the actual message array, so compression (keeping 2 recent screenshots, 1 recent ariaTree) now actually reduces memory usage.

Confidence Score: 5/5

  • This PR is safe to merge with high confidence - it fixes a real memory leak issue with a clean refactor
  • The changes are well-architected and fix a genuine bug where compression wasn't reducing memory. The refactor is surgical: moving compression from middleware (which copied data) to prepareStep (which mutates in-place). The logic correctly handles both output.value and result fields. Existing tests verify message continuation works. No breaking changes to the public API.
  • No files require special attention

Important Files Changed

Filename Overview
packages/core/lib/v3/agent/utils/messageProcessing.ts Refactored to mutate messages in-place for effective memory compression; removed dead code and simplified compression logic
packages/core/lib/v3/handlers/v3AgentHandler.ts Moved message compression from middleware to prepareStep callback, enabling in-place mutation for actual memory reduction

Sequence Diagram

sequenceDiagram
    participant Agent as V3AgentHandler
    participant PrepareStep as prepareStep Callback
    participant LLM as Language Model
    participant Messages as Message Array
    
    Note over Agent,Messages: Before: Middleware Approach
    Agent->>LLM: Call with middleware
    LLM->>Messages: Copy messages
    Messages-->>LLM: Return copied messages
    Note over Messages: Original messages unchanged<br/>Memory not reduced
    
    Note over Agent,Messages: After: PrepareStep Approach
    Agent->>PrepareStep: createPrepareStep()
    PrepareStep->>Messages: processMessages(options.messages)
    Messages->>Messages: Mutate in-place<br/>Replace base64 with placeholders
    Note over Messages: Keep 2 recent screenshots<br/>Keep 1 recent ariaTree
    Messages-->>PrepareStep: Return mutated messages
    PrepareStep->>LLM: Pass to model
    Note over Messages: Memory actually reduced
Loading

Copy link
Contributor

@greptile-apps greptile-apps bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 files reviewed, 1 comment

Edit Code Review Agent Settings | Greptile

userCallback?: PrepareStepFunction<ToolSet>,
): PrepareStepFunction<ToolSet> {
return async (options) => {
processMessages(options.messages);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

style: Consider using return value from processMessages() for logging/debugging compressed message count, or remove return value if not needed

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/core/lib/v3/handlers/v3AgentHandler.ts
Line: 135:135

Comment:
**style:** Consider using return value from `processMessages()` for logging/debugging compressed message count, or remove return value if not needed

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants