rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rename CheckpointStatus to AgentStatus

Browse Source 
The status values describe the agent state, not the
checkpoint data state.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
Aurélien Sibiril
2026-04-13 23:01:37 +02:00
parent ff18a5fc23
commit 71e2d00b3a
7 changed files with 1060 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
pkg/coredata/agent_run.go
View File

@@ -419,7 +419,7 @@ RETURNING
// ClearCheckpoint is the explicit path for removing persisted checkpoint
// data. AgentRun.Update intentionally does not write checkpoint so status
// commits cannot erase a checkpoint saved by PGCheckpointStore.Save.
// commits cannot erase a checkpoint saved by PGCheckpointer.Save.
func (e *AgentRun) ClearCheckpoint(
ctx context.Context,
tx pg.Tx,
@@ -651,20 +651,20 @@ func LoadRunningStopRequestedIDs(ctx context.Context, conn pg.Querier) ([]string
return ids, nil
}
// PGCheckpointStore implements agent.CheckpointStore backed by the
// PGCheckpointer implements agent.Checkpointer backed by the
// agent_runs table checkpoint column. It is supervisor-internal and
// intentionally uses raw run IDs with no tenant scope; public service/API
// methods must load AgentRun through scoped coredata methods before invoking
// lifecycle transitions.
type PGCheckpointStore struct {
type PGCheckpointer struct {
pg *pg.Client
}
func NewPGCheckpointStore(pgClient *pg.Client) *PGCheckpointStore {
return &PGCheckpointStore{pg: pgClient}
func NewPGCheckpointer(pgClient *pg.Client) *PGCheckpointer {
return &PGCheckpointer{pg: pgClient}
}
func (s *PGCheckpointStore) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error {
func (s *PGCheckpointer) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error {
data, err := marshalAgentCheckpoint(cp)
if err != nil {
return err
@@ -694,7 +694,7 @@ func (s *PGCheckpointStore) Save(ctx context.Context, runID string, cp *agent.Ch
)
}
func (s *PGCheckpointStore) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) {
func (s *PGCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) {
var cp *agent.Checkpoint
err := s.pg.WithConn(
@@ -730,10 +730,6 @@ func (s *PGCheckpointStore) Load(ctx context.Context, runID string) (*agent.Chec
return fmt.Errorf("cannot unmarshal checkpoint: %w", err)
}
if cp.Version != agent.CheckpointVersion {
return fmt.Errorf("cannot load checkpoint: unsupported checkpoint version %d", cp.Version)
}
return nil
},
)
@@ -741,42 +737,12 @@ func (s *PGCheckpointStore) Load(ctx context.Context, runID string) (*agent.Chec
return cp, err
}
func (s *PGCheckpointStore) Delete(ctx context.Context, runID string) error {
return s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
q := `UPDATE agent_runs SET checkpoint = NULL, updated_at = now() WHERE id = @id;`
args := pgx.StrictNamedArgs{"id": runID}
_, err := conn.Exec(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot delete checkpoint: %w", err)
}
// Delete is intentionally idempotent: callers use it for cleanup after
// completion, and a concurrently-cleared checkpoint is already the
// desired state.
return nil
},
)
}
func marshalAgentCheckpoint(cp *agent.Checkpoint) ([]byte, error) {
if cp == nil {
return nil, fmt.Errorf("cannot marshal checkpoint: checkpoint is required")
}
next := *cp
if next.Version == 0 {
next.Version = agent.CheckpointVersion
}
if next.Version != agent.CheckpointVersion {
return nil, fmt.Errorf("cannot marshal checkpoint: unsupported checkpoint version %d", next.Version)
}
data, err := json.Marshal(&next)
data, err := json.Marshal(cp)
if err != nil {
return nil, fmt.Errorf("cannot marshal checkpoint: %w", err)
}

122
pkg/coredata/pg_checkpointer_test.go Normal file
View File

@@ -0,0 +1,122 @@
// 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")
},
)
}