New stream not initiated inside frontend when reopening and executing a cell

[jlewi]

With stateful/vscode-runme#1744 I noticed an edge case where a new GenerateStream was not initiated and instead the existing stream was reused. This lead the a failure because the sessionIds didn't match so the server detects that the update is invalid and that a new stream should have been initiated.

Here's the relevant server logs.

2024-10-17T14:53:10.228-0700	INFO	agent/agent.go:593	LogEvent	{"eventId": "01JAE7GTJF4QT157A10MD3JFN2", "eventType": "SESSION_END", "contextId": "01JAE7FX607M7YY8N63QPVX1H3", "selectedCellId": "", "event": {"contextId":"01JAE7FX607M7YY8N63QPVX1H3","eventId":"01JAE7GTJF4QT157A10MD3JFN2","type":"SESSION_END"}}
[GIN] 2024/10/17 - 14:53:10 | 200 |    1.458458ms |       127.0.0.1 | POST     "/api/AIService/LogEvents"
2024-10-17T14:53:10.230-0700	INFO	agent/agent.go:593	LogEvent	{"eventId": "01JAE7GTJFZ0CKRZ25Z2YQM374", "eventType": "SESSION_START", "contextId": "01JAE7GTJE1W81BMDGBQ9W61JA", "selectedCellId": "", "event": {"contextId":"01JAE7GTJE1W81BMDGBQ9W61JA","eventId":"01JAE7GTJFZ0CKRZ25Z2YQM374","type":"SESSION_START"}}
[GIN] 2024/10/17 - 14:53:10 | 200 |     3.60875ms |       127.0.0.1 | POST     "/api/AIService/LogEvents"
2024-10-17T14:53:12.406-0700	INFO	agent/agent.go:411	Request is invalid	{"traceId": "9ed5dd7e7b0dd1d9c7eb5c99bdf14d6c", "evalMode": false, "err": "rpc error: code = InvalidArgument desc = ContextId doesn't match current value; expected 01JAE7FX607M7YY8N63QPVX1H3; got 01JAE7GTJE1W81BMDGBQ9W61JA"}
2024-10-17T14:53:12.407-0700	INFO	agent/agent.go:292	Context cancelled; stopping completion generation	{"traceId": "9ed5dd7e7b0dd1d9c7eb5c99bdf14d6c", "evalMode": false}
[GIN] 2024/10/17 - 14:53:12 | 200 | 27.583013125s |       127.0.0.1 | POST     "/api/AIService/StreamGenerate"
2024-10-17T14:53:12.408-0700	INFO	agent/agent.go:593	LogEvent	{"eventId": "01JAE7GWPKFJ6K5D7KQFDVHAW1", "eventType": "EXECUTE", "contextId": "01JAE7GTJE1W81BMDGBQ9W61JA", "selectedCellId": "01JAE77XFX3Z9H1GQBMVKKQFDF", "event": {"cells":[{"kind":"CELL_KIND_CODE","metadata":{"interactive":"false","runme.dev/id":"01JAE77XFX3Z9H1GQBMVKKQFDF"},"value":"echo \"Fi fi fo\""}],"contextId":"01JAE7GTJE1W81BMDGBQ9W61JA","eventId":"01JAE7GWPKFJ6K5D7KQFDVHAW1","selectedId":"01JAE77XFX3Z9H1GQBMVKKQFDF","type":"EXECUTE"}}

Here's what I believe to be the root cause.
In the VSCode extension, the GhostCellGenerator keeps track of the active cell that is associated with the current stream. It uses this to detect when the active cell changes and therefore a new stream should be initiated. It does this by storing the activeCell in a map keyed by notebookUri
https://github.com/stateful/vscode-runme/blob/90fcc5b60322b4fe34be714ac120a5794004f837/src/extension/ai/ghost.ts#L40

I think the problem is when we close the notebook we don't clear the state in the map. Therefore there is an edge case where you

1. Close a document
2. Reopen a document

If the active cell on initial open is the same as the one that was active on close, then the ghost cell generator assumes the existing stream is for the current cell and a new one doesn't need to be initiated.

Since the server detects that the update is invalid and closes the stream, the client will initiate a new stream on the next cellChangeEvent. If the user is editing the cell then a new event will be generated and thus a new stream is generated.
However, if the user is just executing a cell then there is no follow up event to initiate a new stream.

There's two potential ways to fix this.

1. We could try to clear the nbstate when the document closes
2. We could add retries to the stream generation code

2 might be better solution because if a user expects a completion and this doesn't happen due to a transient issue (i.e. this issue or some other error) then we'd like to retry automatically because not generating a completion is frustrating.