71e2d00b3a
The status values describe the agent state, not the checkpoint data state. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
123 lines
3.6 KiB
Go
123 lines
3.6 KiB
Go
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
|
//
|
|
// Permission to use, copy, modify, and/or distribute this software for any
|
|
// purpose with or without fee is hereby granted, provided that the above
|
|
// copyright notice and this permission notice appear in all copies.
|
|
//
|
|
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
|
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
|
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
|
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
|
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
|
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
|
// PERFORMANCE OF THIS SOFTWARE.
|
|
|
|
package coredata_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.probo.inc/probo/pkg/agent"
|
|
"go.probo.inc/probo/pkg/agentruntest"
|
|
"go.probo.inc/probo/pkg/coredata"
|
|
"go.probo.inc/probo/pkg/llm"
|
|
)
|
|
|
|
func TestPGCheckpointer(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
client := agentruntest.PGClient(t)
|
|
store := coredata.NewPGCheckpointer(client)
|
|
ctx := context.Background()
|
|
|
|
run := agentruntest.InsertPendingRun(
|
|
t,
|
|
client,
|
|
"test-agent",
|
|
[]llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "hello"}}}},
|
|
)
|
|
runID := run.ID.String()
|
|
|
|
t.Run(
|
|
"load returns nil when no checkpoint exists",
|
|
func(t *testing.T) {
|
|
cp, err := store.Load(ctx, runID)
|
|
require.NoError(t, err)
|
|
assert.Nil(t, cp)
|
|
},
|
|
)
|
|
|
|
t.Run(
|
|
"save and load round-trip",
|
|
func(t *testing.T) {
|
|
original := &agent.Checkpoint{
|
|
Status: agent.AgentStatusSuspended,
|
|
AgentName: "test-agent",
|
|
Messages: []llm.Message{
|
|
{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "hello"}}},
|
|
{Role: llm.RoleAssistant, Parts: []llm.Part{llm.TextPart{Text: "working..."}}},
|
|
},
|
|
Usage: llm.Usage{InputTokens: 20, OutputTokens: 10},
|
|
Turns: 1,
|
|
ToolUsedInRun: true,
|
|
}
|
|
|
|
err := store.Save(ctx, runID, original)
|
|
require.NoError(t, err)
|
|
|
|
loaded, err := store.Load(ctx, runID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, loaded)
|
|
|
|
assert.Equal(t, original.Status, loaded.Status)
|
|
assert.Equal(t, original.AgentName, loaded.AgentName)
|
|
assert.Equal(t, original.Usage, loaded.Usage)
|
|
assert.Equal(t, original.Turns, loaded.Turns)
|
|
assert.Equal(t, original.ToolUsedInRun, loaded.ToolUsedInRun)
|
|
assert.Len(t, loaded.Messages, 2)
|
|
},
|
|
)
|
|
|
|
t.Run(
|
|
"save overwrites previous checkpoint",
|
|
func(t *testing.T) {
|
|
updated := &agent.Checkpoint{
|
|
Status: agent.AgentStatusSuspended,
|
|
AgentName: "test-agent",
|
|
Messages: []llm.Message{
|
|
{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "hello"}}},
|
|
{Role: llm.RoleAssistant, Parts: []llm.Part{llm.TextPart{Text: "working..."}}},
|
|
{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "continue"}}},
|
|
},
|
|
Usage: llm.Usage{InputTokens: 30, OutputTokens: 15},
|
|
Turns: 2,
|
|
}
|
|
|
|
err := store.Save(ctx, runID, updated)
|
|
require.NoError(t, err)
|
|
|
|
loaded, err := store.Load(ctx, runID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, loaded)
|
|
assert.Equal(t, 2, loaded.Turns)
|
|
assert.Len(t, loaded.Messages, 3)
|
|
},
|
|
)
|
|
|
|
t.Run(
|
|
"save to nonexistent run returns error",
|
|
func(t *testing.T) {
|
|
cp := &agent.Checkpoint{
|
|
Status: agent.AgentStatusSuspended,
|
|
AgentName: "test-agent",
|
|
}
|
|
err := store.Save(ctx, "nonexistent-run-id", cp)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "not found")
|
|
},
|
|
)
|
|
}
|