Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 104 additions & 0 deletions src/__tests__/store.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,110 @@ describe("Store", () => {

expect(getState().isProcessing).toBe(false)
})

test("abortRequest should save partially streamed message to chat history", () => {
const { getState, actions } = setup()

const controller = new AbortController()
actions.setAbortController(controller)
actions.setIsProcessing(true)
actions.setCurrentAssistantMessage("This is a partial response...")

actions.abortRequest()

const state = getState()

// Message should be saved to chat history
expect(state.chatHistory).toHaveLength(1)
expect(state.chatHistory[0]).toEqual({
type: "message",
role: "assistant",
content: "This is a partial response...",
})

// Current message should be cleared
expect(state.currentAssistantMessage).toBe("")

// Stats should indicate abort
expect(state.stats).toBe("User aborted the request.")

// Processing should be stopped
expect(state.isProcessing).toBe(false)
})

test("abortRequest should not save empty message to chat history", () => {
const { getState, actions } = setup()

const controller = new AbortController()
actions.setAbortController(controller)
actions.setIsProcessing(true)
actions.setCurrentAssistantMessage("")

actions.abortRequest()

const state = getState()

// No message should be added to history
expect(state.chatHistory).toHaveLength(0)

// Current message should still be empty
expect(state.currentAssistantMessage).toBe("")

// Stats should still indicate abort
expect(state.stats).toBe("User aborted the request.")
})

test("abortRequest should not save whitespace-only message to chat history", () => {
const { getState, actions } = setup()

const controller = new AbortController()
actions.setAbortController(controller)
actions.setIsProcessing(true)
actions.setCurrentAssistantMessage(" \n \t ")

actions.abortRequest()

const state = getState()

// No message should be added to history (whitespace is trimmed)
expect(state.chatHistory).toHaveLength(0)

// Current message should still be empty
expect(state.currentAssistantMessage).toBe("")
})

test("abortRequest should preserve existing chat history when adding partial message", () => {
const { getState, actions } = setup()

// Add existing history
actions.addChatHistoryEntry({
type: "message",
role: "user",
content: "Hello",
})

const controller = new AbortController()
actions.setAbortController(controller)
actions.setIsProcessing(true)
actions.setCurrentAssistantMessage("Partial response")

actions.abortRequest()

const state = getState()

// Should have both messages
expect(state.chatHistory).toHaveLength(2)
expect(state.chatHistory[0]).toEqual({
type: "message",
role: "user",
content: "Hello",
})
expect(state.chatHistory[1]).toEqual({
type: "message",
role: "assistant",
content: "Partial response",
})
})
})

describe("reset action", () => {
Expand Down
11 changes: 11 additions & 0 deletions src/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,18 @@ export const AgentStore = createContextStore<StoreModel>({
abortRequest: action((state) => {
state.abortController?.abort()
state.abortController = undefined

if (state.currentAssistantMessage.trim()) {
// Persist what has been generated so far
state.chatHistory.push({
type: "message",
role: "assistant",
content: state.currentAssistantMessage,
})
}

state.currentAssistantMessage = ""

state.stats = "User aborted the request."
state.isProcessing = false
}),
Expand Down