-
Notifications
You must be signed in to change notification settings - Fork 97
/
memops.go
219 lines (196 loc) · 6.92 KB
/
memops.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
// Copyright 2023, Command Line Inc.
// SPDX-License-Identifier: Apache-2.0
// in-memory storage for waveterm server
package sstore
import (
"fmt"
"log"
"sync"
"github.com/wavetermdev/waveterm/waveshell/pkg/packet"
"github.com/wavetermdev/waveterm/waveshell/pkg/utilfn"
)
// global lock for all memory operations
// memory ops are very fast, so this is not a bottleneck
var MemLock *sync.Mutex = &sync.Mutex{}
var ScreenMemStore map[string]*ScreenMemState = make(map[string]*ScreenMemState) // map of screenid -> ScreenMemState
type StatusIndicatorLevel int
const (
StatusIndicatorLevel_None StatusIndicatorLevel = iota
StatusIndicatorLevel_Output
StatusIndicatorLevel_Success
StatusIndicatorLevel_Error
)
func dumpScreenMemStore() {
MemLock.Lock()
defer MemLock.Unlock()
for k, v := range ScreenMemStore {
log.Printf(" ScreenMemStore[%s] = %+v\n", k, v)
}
}
type OpenAICmdInfoChatStore struct {
MessageCount int `json:"messagecount"`
Messages []*packet.OpenAICmdInfoChatMessage `json:"messages"`
}
type ScreenMemState struct {
NumRunningCommands int `json:"numrunningcommands,omitempty"`
StatusIndicator StatusIndicatorLevel `json:"statusindicator,omitempty"`
CmdInputText utilfn.StrWithPos `json:"cmdinputtext,omitempty"`
CmdInputSeqNum int `json:"cmdinputseqnum,omitempty"`
AICmdInfoChat *OpenAICmdInfoChatStore `json:"aicmdinfochat,omitempty"`
}
func ScreenMemDeepCopyCmdInfoChatStore(store *OpenAICmdInfoChatStore) *OpenAICmdInfoChatStore {
rtnMessages := []*packet.OpenAICmdInfoChatMessage{}
for index := 0; index < len(store.Messages); index++ {
messageToCopy := *store.Messages[index]
if messageToCopy.AssistantResponse != nil {
assistantResponseCopy := *messageToCopy.AssistantResponse
messageToCopy.AssistantResponse = &assistantResponseCopy
}
rtnMessages = append(rtnMessages, &messageToCopy)
}
rtn := &OpenAICmdInfoChatStore{MessageCount: store.MessageCount, Messages: rtnMessages}
return rtn
}
func ScreenMemInitCmdInfoChat(screenId string) {
greetingMessagePk := &packet.OpenAICmdInfoChatMessage{
MessageID: 0,
IsAssistantResponse: true,
AssistantResponse: &packet.OpenAICmdInfoPacketOutputType{
Message: packet.OpenAICmdInfoChatGreetingMessage,
},
}
ScreenMemStore[screenId].AICmdInfoChat = &OpenAICmdInfoChatStore{MessageCount: 1, Messages: []*packet.OpenAICmdInfoChatMessage{greetingMessagePk}}
}
func ScreenMemClearCmdInfoChat(screenId string) {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
ScreenMemInitCmdInfoChat(screenId)
}
func ScreenMemAddCmdInfoChatMessage(screenId string, msg *packet.OpenAICmdInfoChatMessage) {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
if ScreenMemStore[screenId].AICmdInfoChat == nil {
log.Printf("AICmdInfoChat is null, creating")
ScreenMemInitCmdInfoChat(screenId)
}
CmdInfoChat := ScreenMemStore[screenId].AICmdInfoChat
CmdInfoChat.Messages = append(CmdInfoChat.Messages, msg)
CmdInfoChat.MessageCount++
}
func ScreenMemGetCmdInfoMessageCount(screenId string) int {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
if ScreenMemStore[screenId].AICmdInfoChat == nil {
ScreenMemInitCmdInfoChat(screenId)
}
return ScreenMemStore[screenId].AICmdInfoChat.MessageCount
}
func ScreenMemGetCmdInfoChat(screenId string) *OpenAICmdInfoChatStore {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
if ScreenMemStore[screenId].AICmdInfoChat == nil {
ScreenMemInitCmdInfoChat(screenId)
}
return ScreenMemDeepCopyCmdInfoChatStore(ScreenMemStore[screenId].AICmdInfoChat)
}
func ScreenMemUpdateCmdInfoChatMessage(screenId string, messageID int, msg *packet.OpenAICmdInfoChatMessage) error {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
if ScreenMemStore[screenId].AICmdInfoChat == nil {
ScreenMemInitCmdInfoChat(screenId)
}
CmdInfoChat := ScreenMemStore[screenId].AICmdInfoChat
if messageID >= 0 && messageID < len(CmdInfoChat.Messages) {
CmdInfoChat.Messages[messageID] = msg
} else {
return fmt.Errorf("ScreenMemUpdateCmdInfoChatMessage: error: Message Id out of range: %d", messageID)
}
return nil
}
func ScreenMemSetCmdInputText(screenId string, sp utilfn.StrWithPos, seqNum int) {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
if seqNum <= ScreenMemStore[screenId].CmdInputSeqNum {
return
}
ScreenMemStore[screenId].CmdInputText = sp
ScreenMemStore[screenId].CmdInputSeqNum = seqNum
}
func ScreenMemIncrementNumRunningCommands(screenId string, delta int) int {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
newNum := ScreenMemStore[screenId].NumRunningCommands + delta
ScreenMemStore[screenId].NumRunningCommands = newNum
return newNum
}
// If the new indicator level is higher than the current indicator, update the current indicator. Returns the new indicator level.
func ScreenMemCombineIndicatorLevels(screenId string, level StatusIndicatorLevel) StatusIndicatorLevel {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
curLevel := ScreenMemStore[screenId].StatusIndicator
if level > curLevel {
ScreenMemStore[screenId].StatusIndicator = level
return level
} else {
return curLevel
}
}
// Set the indicator to the given level, regardless of the current indicator level.
func ScreenMemSetIndicatorLevel(screenId string, level StatusIndicatorLevel) {
MemLock.Lock()
defer MemLock.Unlock()
if ScreenMemStore[screenId] == nil {
ScreenMemStore[screenId] = &ScreenMemState{}
}
ScreenMemStore[screenId].StatusIndicator = StatusIndicatorLevel_None
}
func GetCurrentIndicatorState() ([]*ScreenStatusIndicatorType, []*ScreenNumRunningCommandsType) {
MemLock.Lock()
defer MemLock.Unlock()
indicators := []*ScreenStatusIndicatorType{}
numRunningCommands := []*ScreenNumRunningCommandsType{}
for screenId, screenMem := range ScreenMemStore {
if screenMem.StatusIndicator > 0 {
indicators = append(indicators, &ScreenStatusIndicatorType{ScreenId: screenId, Status: screenMem.StatusIndicator})
}
if screenMem.NumRunningCommands > 0 {
numRunningCommands = append(numRunningCommands, &ScreenNumRunningCommandsType{ScreenId: screenId, Num: screenMem.NumRunningCommands})
}
}
return indicators, numRunningCommands
}
// safe because we return a copy
func GetScreenMemState(screenId string) *ScreenMemState {
MemLock.Lock()
defer MemLock.Unlock()
ptr := ScreenMemStore[screenId]
if ptr == nil {
return nil
}
rtn := *ptr
return &rtn
}