8000 feat(run): add SessionRunStore adapter by Fodesu · Pull Request #85 · felinics/twilight · GitHub
[go: up one dir, main page]

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
55 changes: 55 additions & 0 deletions agentcore/run/sessionstore/cachepolicy_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
package sessionstore

import (
"github.com/felinics/twilight/agentcore/ledger"
"github.com/felinics/twilight/agentcore/module"
"testing"
)

// WriterCachePolicy is the module's answer to "who may cache this projection":
// the Writer may cache every projection except the machine projection, whose
// checkpoints the Runtime owns (RUN-CMT-2), and the interval is the
// deployment's, not a constant of this module.
func TestWriterCachePolicyExcludesTheMachineProjection(t *testing.T) {
other := module.ProjectionID("twilight/chatlog/surface")
for _, every := range []ledger.CommitSeq{1, module.DefaultCacheEvery, 4096} {
policy := WriterCachePolicy(every)
// The interval has elapsed for any of these, so only the exclusion can
// decline.
head := ledger.Head{Next: every + 1}
if policy(MachineProjectionID, MachineProjection.Version, head, ledger.Head{}, false) {
t.Errorf("every=%d: the Writer would cache the machine projection", every)
}
if policy(MachineProjectionID, MachineProjection.Version, head, ledger.Head{}, true) {
t.Errorf("every=%d: the Writer would cache the machine projection at Close", every)
}
if !policy(other, 1, head, ledger.Head{}, false) {
t.Errorf("every=%d: another projection is not cached once the interval elapsed", every)
}
}
}

// The interval is honoured, including the zero value standing in for the
// default, so a deployment that passes no interval still gets bounded work.
func TestWriterCachePolicyHonoursTheInterval(t *testing.T) {
other := module.ProjectionID("twilight/chatlog/surface")
cases := []struct {
every ledger.CommitSeq
head ledger.CommitSeq
cached ledger.CommitSeq
want bool
}{
{every: 8, head: 7, cached: 0, want: false},
{every: 8, head: 8, cached: 0, want: true},
{every: 8, head: 12, cached: 5, want: false},
{every: 8, head: 13, cached: 5, want: true},
{every: 0, head: module.DefaultCacheEvery - 1, cached: 0, want: false},
{every: 0, head: module.DefaultCacheEvery, cached: 0, want: true},
}
for _, tc := range cases {
got := WriterCachePolicy(tc.every)(other, 1, ledger.Head{Next: tc.head}, ledger.Head{Next: tc.cached}, false)
if got != tc.want {
t.Errorf("every=%d head=%d cached=%d: got %v, want %v", tc.every, tc.head, tc.cached, got, tc.want)
}
}
}
93 changes: 93 additions & 0 deletions agentcore/run/sessionstore/codec_history_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
package sessionstore

import (
"errors"
"github.com/felinics/twilight/agentcore/jsonstable"
"github.com/felinics/twilight/agentcore/ledger"
"github.com/felinics/twilight/agentcore/module"
"github.com/felinics/twilight/agentcore/run"
"strings"
"testing"
)

// withdrawnV1 stands in for a superseded wire shape of model_step_withdrawn
// in which the step was named "step": it decodes that shape and upcasts to
// the current run.ModelStepWithdrawn. Encode is never used for a superseded
// version.
type withdrawnV1 struct{}

func (withdrawnV1) Validate(any) error { return errors.New("version 1 is superseded") }
func (withdrawnV1) Encode(any) (jsonstable.Value, error) {
return jsonstable.Value{}, errors.New("version 1 is superseded")
}
func (withdrawnV1) Decode(w jsonstable.Value) (any, error) {
var body struct {
RunID run.RunID `json:"runId"`
Step run.StepID `json:"step"`
}
if err := w.Decode(&body); err != nil {
return nil, err
}
return Event{RunID: body.RunID, Fact: run.ModelStepWithdrawn{StepID: body.Step}}, nil
}

// EXT-REG-2: a fact type keeps the codec of every payload version it ever
// wrote. With a superseded version in its history the type encodes at the
// next version and still decodes the old wire to the current fact.
func TestFactCodecHistoryDecodesEveryVersion(t *testing.T) {
const name = "model_step_withdrawn"
codecs, current := factCodecs(name, map[module.PayloadVersion]module.PayloadCodec{1: withdrawnV1{}})
if current != 2 || len(codecs) != 2 {
t.Fatalf("history = %d codecs at version %d, want 2 at 2", len(codecs), current)
}
mod := module.ModuleDescriptor{Source: module.SourceTwilight, ID: ModuleID,
Streams: []module.StreamDefinition{streamDefinition}, Events: []module.EventDefinition{eventDefinition(name, codecs, current)}}
reg, err := module.BuildRegistry(mod)
if err != nil {
t.Fatal(err)
}
typ := Type(name)
want := Event{RunID: "r1", Fact: run.ModelStepWithdrawn{StepID: "s1"}}

wire, err := reg.Encode(typ, want)
if err != nil || !strings.Contains(wire.String(), `"v":2`) {
t.Fatalf("encode = %s %v, want the current version 2", wire, err)
}
cases := []struct {
name string
payload string
version module.PayloadVersion
}{
{"superseded wire", `{"runId":"r1","step":"s1","v":1}`, 1},
{"current wire", wire.String(), 2},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
d, err := reg.Decode(ledger.Event{Type: typ, Payload: jsonstable.MustParse(tc.payload)})
if err != nil || d.Unknown || d.Version != tc.version || d.Value != want {
t.Fatalf("decode = %+v %v, want version %d value %+v", d, err, tc.version, want)
}
})
}
}

// The production module has no history yet: every fact type is at version 1
// with one codec. A fact whose wire shape changes must add its previous
// codec to olderFactCodecs rather than edit the current one in place.
func TestEveryFactIsAtVersionOne(t *testing.T) {
for _, def := range Module.Events {
if def.Version != 1 || len(def.Codecs) != 1 || def.Codecs[1] == nil {
t.Fatalf("%s: version %d with %d codecs, want version 1 with one codec", def.Type, def.Version, len(def.Codecs))
}
}
}

// A history with a gap is caught when the module is built.
func TestFactCodecHistoryMustBeContiguous(t *testing.T) {
defer func() {
if recover() == nil {
t.Fatal("a history without version 1 built")
}
}()
factCodecs("model_step_withdrawn", map[module.PayloadVersion]module.PayloadCodec{2: withdrawnV1{}})
}
226 changes: 226 additions & 0 deletions agentcore/run/sessionstore/frozen.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,226 @@
package sessionstore

import (
"bytes"
"context"
"errors"
"fmt"
"io"
"strings"

"github.com/felinics/twilight/agentcore/artifact"
"github.com/felinics/twilight/agentcore/run"
"github.com/felinics/twilight/agentcore/run/frozen"
"github.com/felinics/twilight/agentcore/run/model"
"github.com/felinics/twilight/agentcore/run/store"
"github.com/felinics/twilight/agentcore/run/wire"
)

// FrozenAuthority is the cas Authority of frozen run bodies: the run layer's
// frozen.Store is one Authority of the artifact ContentStore, not a
// second content-addressed store (RUN-WIR-4).
const FrozenAuthority artifact.Authority = frozen.Authority

// FrozenMediaType is the media type frozen bodies are stored under; the cas
// store binds it to the key, so every Put and Get of a body agrees on it.
const FrozenMediaType = "application/vnd.twilight.frozen+json"

// frozenValues realizes frozen.Store over a cas ContentStore and the
// BindingStore the Session Writers admit against. A body's digest is the
// SHA-256 of its bytes (run.EncodeFrozen*), which is exactly the cas Key, so
// no index maps digests to refs: Get rebuilds the Ref from the digest alone.
// Put stores the body and registers its Binding in one call, so a fact that
// names the digest is admissible as soon as Put returns and no caller has to
// coordinate two stores before a commit (RUN-WIR-4, EXT-WRT-3). Bodies are
// EventBound content whose retention root is the fact that names them; the
// adapter offers no Delete.
type frozenValues struct {
store artifact.ContentStore
bindings artifact.BindingStore
}

// FrozenValues adapts a cas ContentStore serving FrozenAuthority and the
// Writers' BindingStore to the run layer's frozen.Store port. bindings
// must be the store the Writers' Admission resolves against. A content store
// of another Authority answers every Get with ErrUnauthorized, which surfaces
// as an error rather than a miss.
func FrozenValues(content artifact.ContentStore, bindings artifact.BindingStore) (frozen.Store, error) {
if content == nil || bindings == nil {
return nil, errors.New("sessionstore: frozen values require a content store and a binding store")
}
return &frozenValues{store: content, bindings: bindings}, nil
}

func (f *frozenValues) Put(ctx context.Context, digest run.Digest, val []byte) error {
if err := ctx.Err(); err != nil {
return err
}
ref, err := FrozenRef(digest)
if err != nil {
return err
}
if key, _ := artifact.CASKey(val); key != ref.Key {
return fmt.Errorf("agent: frozen values: body digests to %s, not to its name %s", key, digest)
}
if _, err := f.store.Put(ctx, artifact.PutRequest{MediaType: FrozenMediaType, Reader: bytes.NewReader(val), Durability: artifact.EventBound}); err != nil {
return err
}
binding, err := FrozenBinding(digest)
if err != nil {
return err
}
// An identical Binding is already registered on a replay; only a
// differing Ref under the same BindingID conflicts (ART-BND-1).
_, err = f.bindings.CreateBinding(ctx, binding)
return err
}

func (f *frozenValues) Get(ctx context.Context, digest run.Digest) (body []byte, found bool, err error) {
if err := ctx.Err(); err != nil {
return nil, false, err
}
if digest == "" {
return nil, false, nil
}
ref, err := FrozenRef(digest)
if err != nil {
return nil, false, err
}
rc, _, err := f.store.Open(ctx, ref)
if err != nil {
var aerr *artifact.Error
if errors.As(err, &aerr) && aerr.Code == artifact.ErrNotFound {
return nil, false, nil
}
return nil, false, err
}
defer rc.Close()
raw, err := io.ReadAll(rc)
if err != nil {
return nil, false, err
}
return raw, true, nil
}

// FrozenRef is the cas Ref of the body named by digest: Key and Integrity are
// both the digest, MediaType and Durability are the adapter's constants. The
// Run protocol knows only the digest; this derivation is the adapter's
// (RUN-WIR-4).
func FrozenRef(digest run.Digest) (artifact.Ref, error) {
algorithm, val, ok := strings.Cut(string(digest), ":")
if !ok || algorithm != artifact.IntegritySHA256 || val == "" {
return artifact.Ref{}, fmt.Errorf("agent: frozen values: digest %q is not sha256:<hex>", digest)
}
return artifact.Ref{Scheme: artifact.SchemeCAS, Authority: FrozenAuthority, Key: artifact.Key(digest), MediaType: FrozenMediaType,
Integrity: &artifact.Integrity{Algorithm: algorithm, Value: val}, Durability: artifact.EventBound}, nil
}

// FrozenBindingID is the BindingID under which a frozen body's Ref is
// registered: derived from the digest alone, so the fact that names the body
// and the binding admission agree without an index.
func FrozenBindingID(digest run.Digest) artifact.BindingID {
return artifact.BindingID(frozen.Authority + "/" + string(digest))
}

// FrozenBinding is the Binding a fact naming digest references. Its Ref is
// FrozenRef(digest), so registering it is idempotent across replays.
func FrozenBinding(digest run.Digest) (artifact.Binding, error) {
ref, err := FrozenRef(digest)
if err != nil {
return artifact.Binding{}, err
}
return artifact.NewBinding(FrozenBindingID(digest), ref)
}

// frozenRefs is the BindingExtractor of the fact types that name a frozen
// body: the one BindingID the Writer admits and claims for the commit
// (EXT-REF-2, EXT-WRT-3).
func frozenRefs(val any) ([]artifact.BindingID, error) {
ev, ok := val.(Event)
if !ok {
return nil, fmt.Errorf("frozen refs: unexpected %T", val)
}
var digest run.Digest
switch f := ev.Fact.(type) {
case run.ModelStepPrepared:
digest = f.RequestDigest
case run.ModelStepCompleted:
digest = f.ResultDigest
case run.ToolCallCompleted:
digest = f.OutputDigest
case run.ToolCallAnswered:
digest = f.ResponseDigest
default:
return nil, fmt.Errorf("frozen refs: %s names no frozen body", wire.FactType(ev.Fact))
}
if digest == "" {
return nil, fmt.Errorf("frozen refs: %s has an empty digest", wire.FactType(ev.Fact))
}
return []artifact.BindingID{FrozenBindingID(digest)}, nil
}

// --- materialization ---------------------------------------------------------------

// Content resolves the bodies facts name by digest (RUN-WIR-4) for the layer
// that renders a structural projection into a usable representation. It is
// the materialization port: projections stay pure folds over facts and never
// touch it.
type Content struct {
Frozen frozen.Store
}

// NewContent builds the materializer over a frozen.Store.
func NewContent(st frozen.Store) *Content { return &Content{Frozen: st} }

func (c *Content) raw(ctx context.Context, what string, digest run.Digest) ([]byte, error) {
if err := store.CheckContext(ctx); err != nil {
return nil, err
}
if digest == "" {
return nil, fmt.Errorf("sessionstore: empty %s digest", what)
}
raw, ok, err := c.Frozen.Get(ctx, digest)
if err != nil {
return nil, err
}
if !ok {
return nil, fmt.Errorf("%w: %s %s", frozen.ErrMissing, what, digest)
}
return raw, nil
}

// ModelResult returns the body ModelStepCompleted.ResultDigest names.
func (c *Content) ModelResult(ctx context.Context, digest run.Digest) (model.ModelResult, error) {
raw, err := c.raw(ctx, "model result", digest)
if err != nil {
return model.ModelResult{}, err
}
return frozen.DecodeModelResult(raw, digest)
}

// ToolOutput returns the body ToolCallCompleted.OutputDigest names.
func (c *Content) ToolOutput(ctx context.Context, digest run.Digest) (run.CanonicalJSON, error) {
raw, err := c.raw(ctx, "tool output", digest)
if err != nil {
return run.CanonicalJSON{}, err
}
return frozen.DecodeToolOutput(raw, digest)
}

// ToolResponse returns the body ToolCallAnswered.ResponseDigest names.
func (c *Content) ToolResponse(ctx context.Context, digest run.Digest) (run.CanonicalJSON, error) {
raw, err := c.raw(ctx, "tool response", digest)
if err != nil {
return run.CanonicalJSON{}, err
}
return frozen.DecodeToolResponse(raw, digest)
}

// ModelRequest returns the body ModelStepPrepared.RequestDigest names.
func (c *Content) ModelRequest(ctx context.Context, digest run.Digest) (model.ModelRequest, error) {
raw, err := c.raw(ctx, "model request", digest)
if err != nil {
return model.ModelRequest{}, err
}
return frozen.DecodeRequest(raw, digest)
}
Loading
Loading
0