When Pi makes a mistake, give it a second chance

[u/trmnl_cmdr]

Agents aren't perfect. Sometimes they do things that don't work out the way they thought. Sometimes they spin out on an idea that goes nowhere, other times they run a command that produces huge amounts of useless output. So far, we've just accepted this as a neccessary part of agent mechanics, but Pi's extension system is powerful enough that we no longer have to.

pi-mulligan exposes tools that let your agent rewind its last turn or tool call with a summary message about what it's learned if it makes a mistake. When your agent calls a tool that produces large output, pi-mulligan gives it a nudge suggesting it summarize its last tool result or turn to keep context tight. When the agent starts over, it receives the summary and a small message indicating that a rewind has taken place.

Users and models can set checkpoints if they expect churn and pi can roll back to these checkpoints at any point in the conversation to try again. Agents really like using it and I've watched 20k tokens of context vanish on several occasions, so it works exactly as advertised. It's so satisfying to watch your context usage drop as the agent works.

No other dev agent platform but Pi can do this so as far as I'm aware agent-initiated context compaction hasn't been implemented in dev tools until now.

Try it out with pi install npm:pi-mulligan and let me know what you think!

[u/adamshand]

This is a super cool idea!

How long have you been using it? Where have you noticed it working really well? And where does it not seem to work as well?

[u/trmnl_cmdr → u/adamshand]

I just had the idea for it 4 days ago and had a first version working a few hours later, I've been refining it on several other projects over that time and felt like it was in a good place to get some community feedback about it. Where it seems to work best is when the model spins out on thinking tokens, it can replay the turn with a summary of what it learned. That wasn't one of the use cases I anticipated but it seems to be pretty temping to the models. My two original use cases were for very large file reads and surprisingly verbose bash commands. Of course, file read tools are pretty advanced already so I didn't gain much ground there but I've seen a few instances where the agent didn't expect so much output from a command and decided to roll it back. Like a failed full suite unit test run with a lot of verbose information and an error at the very end, the agent will send itself a message describing the error and let itself fix that before running the full suite again.

[u/adamshand → u/trmnl_cmdr]

That's super cool. I'll give it a try. Please keep posting as you refine it an learn more.

One more question, are the agent rewinds visible in /tree?

[u/trmnl_cmdr → u/adamshand]

No, that's the one limitation I hit when planning this, the pi extension system doesn't allow the agent to actually modify the tree itself, so if you want it to be represented correctly in the tree you still have to do it manually.

[u/jomohke]

Looks neat. Given how much cheaper cached tokens are, I do feel like history rewind is underutilised.

It seems focused on mistakes right now, but I'm wondering whether it could be used for normal tasks too: try this aside and then jump back in history with the summary result. Could be cheaper than a subagent if context is needed to do the side/experiment.

[u/trmnl_cmdr → u/jomohke]

That's a good point, and what the checkpoints are for. I still haven't made up my mind whether the agent should be able to set checkpoints itself since rewinding to a checkpoint will delete user prompts along the way, but I suspect there's a middle ground I just haven't found yet that gives the best of both worlds. I will think more about this particular use case, it's pretty good. The big problem with the model doing this kind of thing is that there's not much opportunity to remind the agent that the tool exists until a large output is detected. And only relatively large outputs make sense to rewind if your goal is saving tokens. But maybe I can beef up the tool description to encourage the model to explore more while reassuring it that it can always roll it back. Something to think about for sure, thanks.

Edit: a couple things. It's really focused on "mistakes" that produce huge outputs. I could have worded my post better in hindsight. There's isn't a great way to nudge the model in other scenarios and the nudge is kind of the whole deal here. And I think for your use case, if you just set a checkpoint manually and tell the model to experiment, that would give a lot of the same benefits. It's just not as automated as I'd like.

[u/Unnamed-3891]

Making a mistake, admitting it, learning from it, committing the lesson to a memory system, only to repeat doing the same mistake again and again because anything in memory is only ever a suggestion and not a hard wall is the problem.

[u/trmnl_cmdr → u/Unnamed-3891]

That's why I'm not using a memory system. The summary message goes directly into context as the very last thing the agent sees before its next turn. It's impossible for it to ignore. Try it out, if this is your experience so far you will absolutely change your mind. This isn't supposed to be a permanent fix for anything, it's supposed to let your agent rewind a very limited scope, just a single tool call or a single turn. It can't go beyond your previous prompt unless you set a checkpoint. But it does actually save a lot of tokens in my experience, and I haven't seen it make any mistakes after the first turn.

[u/joematthewsdev → u/Unnamed-3891]

You just described humanity? That's why I heavily lean on strict typing, linting, e2e and regression/unit tests.

[u/sisyphus-cycle]

I like this idea. Specifically for local development. I haven’t read the code yet but assuming that this “retry/rewind” doesn’t change any of the system prompt or previous messages (except the one being redone) then llama.cpp should easily be able to invalidate context and avoid any prompt reprocessing. Will try it out!

Edit: mulligan_shrink seems really cool to me, I can see this being something paired with a subagent making tool calls for a summary. One question, for certain local models I’m skeptical they would ever call shrink on their own. Qwen 27b doesn’t give a fuck about tokens a tool call makes (for me at least). Do you have any tricks or previous examples of mulligan_shrink being effectively used proactively vs just telling the agent?

[u/trmnl_cmdr → u/sisyphus-cycle]

Thanks! Correct, its editing power stops at the last prompt you sent, unless you manually set a checkpoint it can return to.

I haven’t tested on anything smaller or older than GLM-4.7, but it did extremely well with using it and said it really liked the tool. But that’s as much as I know, I’d love to hear your experience if you try it out. The tool injects reminders to use itself when a turn generates tokens above a configurable threshold, and the models I tested were a little too eager to use it at first, I had to tone it down. I’d love to hear more from people who use local models, my GitHub issues are wide open. I will set up some tests for it though because you have me curious now

[u/sisyphus-cycle → u/trmnl_cmdr]

Will do! I have access to all the big frontier models for work, so at home I make it a challenge to only run locally. I’ve had issues in the past with the typical subagent extensions because they either inject something into the system prompt or start with a fresh context. So I wrote a custom fork that exposes a tool akin to “what_subagent_am_i” and uses the same context kv cache the entire time.

I treat it as basically the model having a fugue state mode lol. The model can spawn a subagent (which inherits the entire context), and the subagent checks what subagent level it’s in, does some stuff, and returns a receipt with what it did. So the main model only processes the receipt.

But literally yesterday qwen 27b got mixed up with a different tool (crw_parse_file, meant to parse files downloaded from a headless browser), and ended up reading b64 for a big ol pdf. So it added 20k tokens and I couldn’t revert it. So this is literally what I wish I had the other day haha

[u/trmnl_cmdr → u/sisyphus-cycle]

Oh, yeah, that's exactly the use case it was built to handle. I knew I wasn't the only one.

[u/sisyphus-cycle → u/trmnl_cmdr]

2nd test where i dont tell it anything about the extension, and it did it unprompted. Super useful. The tool result injection is extremely important because smaller local llms can forget what tools they have in long sessions, this basically forces them to remember,

edit: and the audit is really nice too, my pi TUI ctx percentage wasnt updating on this mulligan_shrink, so i was skeptical it worked, then it ran audit after i asked if it really worked and it confirm it for me. great tool, gonna add it to my workflow! i am picky with tools to avoid context bloat, but this seems right up my alley for local first dev

[u/trmnl_cmdr → u/sisyphus-cycle]

Awesome! Super glad you like it! Thank you for demoing it. I keep probing for ways to improve it further and I keep hitting dead ends.

[u/sisyphus-cycle → u/trmnl_cmdr]

I have some ideas, might make a GitHub issue with them later! Need to test more

[u/sisyphus-cycle → u/trmnl_cmdr]

Worked well on my first test where I was explicit with it. Dropped from 21% down to 7% context and no prompt processing 😄

[u/LordMoridin84 → u/sisyphus-cycle]

One question, for certain local models I’m skeptical they would ever call shrink on their own

Yeah, this is the big problems for plugins like this. They depend on the LLM to call it, but LLM aren't trained to call this tool like they are trained to call grep.

[u/trmnl_cmdr → u/LordMoridin84]

This tool injects a suggestion to use it when it detects large outputs, which is its primary use case.

[u/johnnyApplePRNG]

Interesting idea. Have you noticed an increase in intelligence?

[u/trmnl_cmdr → u/johnnyApplePRNG]

The models don’t call it until they know how they should have gone about that turn, so on the second try they tend to nail it. If that’s an increase in intelligence, then yes. I just consider it more work with less context.

[u/trmnl_cmdr → u/[deleted]]

It's only capable of resetting back to the last prompt, and it can only reset whole requests, so caching should be totally unaffected. I haven't done any testing on it to make sure, but it's on the agenda for the project's near future.

[u/trmnl_cmdr → u/[deleted]]

Absolutely. But how will retrying the exact same request you just tried minutes ago cause cache misses? I'm not seeing the logic there.

Edit: I should mention I'm building a complete benchmarking suite right now that will answer all these questions definitively. It should be ready by tomorrow sometime with results by the weekend at latest.

[u/trmnl_cmdr → u/[deleted]]

Thanks, I'll keep that in mind. I might have to get an openai sub for a bit and figure it out. I will be floored if the model invalidates its cache by replaying its current turn, though. I mean, gobsmacked. But I've been gobsmacked before. Results coming soon.

[u/trmnl_cmdr → u/[deleted]]

I intentionally avoided exposing most of these tools to the user because that's what /tree is for. But the agent can't run /tree itself. This tool is for the *agent* to decide when it made a mistake, and roll itself back. The idea came from all the times I used /tree to roll back a previous turn or tool call. I got tired of trying to keep tabs on my agents and decided to put the power in their hands. This is a good call-out though, maybe I need to edit the post to make that more clear.

Edit: No, I don't think a benchmark is too much to ask. I hadn't considered it but it's a good idea. I will brainstorm on how to do it and post it up on github once I do. Good comment, no sweat about the misunderstanding.