3 files changed, 63 insertions, 14 deletions
diff --git a/libpanto-go/go.mod b/libpanto-go/go.mod
index 805022f..c8ccb75 100644
--- a/libpanto-go/go.mod
+++ b/libpanto-go/go.mod
@@ -1,3 +1,3 @@
-module github.com/travisjeffery/pantograph/libpanto-go
+module code.tjp.lol/pantograph.git/libpanto-go
 
 go 1.24
diff --git a/libpanto-go/panto.go b/libpanto-go/panto.go
index 786ff1c..5644b61 100644
--- a/libpanto-go/panto.go
+++ b/libpanto-go/panto.go
@@ -202,7 +202,12 @@ type (
 )
 
 type (
-	Agent  struct{ ptr *C.PantoAgent }
+	Agent struct {
+		ptr                 *C.PantoAgent
+		retainedSession     *Session
+		retainedConv        *Conversation
+		retainedStoreCloser func()
+	}
 	Stream struct {
 		ptr  *C.PantoStream
 		err  error
@@ -564,16 +569,7 @@ func (b ContentBlock) Snapshot() OwnedContentBlock {
 }
 
 func NewAgent(cfg Config) (*Agent, error) {
-	st, err := NewNullStore()
-	if err != nil {
-		return nil, err
-	}
-	sess, err := CreateSession(st)
-	st.Close()
-	if err != nil {
-		return nil, err
-	}
-	return NewAgentWithSession(cfg, sess, nil)
+	return NewAgentWithSession(cfg, nil, nil)
 }
 
 func NewAgentWithSession(cfg Config, session *Session, conv *Conversation) (*Agent, error) {
@@ -582,6 +578,11 @@ func NewAgentWithSession(cfg Config, session *Session, conv *Conversation) (*Age
 		return nil, err
 	}
 	defer strings.free()
+	var csession *C.PantoSession
+	if session != nil {
+		csession = session.ptr
+		runtime.SetFinalizer(session, nil)
+	}
 	var cconv *C.PantoConversation
 	if conv != nil {
 		cconv = conv.ptr
@@ -589,10 +590,10 @@ func NewAgentWithSession(cfg Config, session *Session, conv *Conversation) (*Age
 		runtime.SetFinalizer(conv, nil)
 	}
 	var out *C.PantoAgent
-	if C.panto_agent_create(&ccfg, session.ptr, cconv, &out) != C.PANTO_OK {
+	if C.panto_agent_create(&ccfg, csession, cconv, &out) != C.PANTO_OK {
 		return nil, lastError()
 	}
-	a := &Agent{out}
+	a := &Agent{ptr: out, retainedSession: session, retainedConv: conv}
 	runtime.SetFinalizer(a, (*Agent).Close)
 	return a, nil
 }
@@ -602,6 +603,20 @@ func (a *Agent) Close() {
 		C.panto_agent_destroy(a.ptr)
 		a.ptr = nil
 	}
+	if a != nil {
+		if a.retainedSession != nil {
+			a.retainedSession.ptr = nil
+			a.retainedSession = nil
+		}
+		if a.retainedConv != nil {
+			a.retainedConv.ptr = nil
+			a.retainedConv = nil
+		}
+		if a.retainedStoreCloser != nil {
+			a.retainedStoreCloser()
+			a.retainedStoreCloser = nil
+		}
+	}
 }
 func (a *Agent) SessionID() string { return goSlice(C.panto_agent_session_id(a.ptr)) }
 func (a *Agent) Conversation() *Conversation {
diff --git a/libpanto-go/panto_test.go b/libpanto-go/panto_test.go
new file mode 100644
index 0000000..ef21f1e
--- /dev/null
+++ b/libpanto-go/panto_test.go
@@ -0,0 +1,34 @@
+package panto
+
+import "testing"
+
+func testConfig() Config {
+	return Config{
+		Provider: ProviderConfig{
+			OpenAIChat: &OpenAIChatConfig{
+				APIKey:    "test-key",
+				BaseURL:   "http://localhost",
+				Model:     "test-model",
+				MaxTokens: 16,
+			},
+		},
+	}
+}
+
+func TestNewAgentSetSystemPromptDoesNotCrash(t *testing.T) {
+	Init()
+	defer Deinit()
+
+	a, err := NewAgent(testConfig())
+	if err != nil {
+		t.Fatalf("NewAgent: %v", err)
+	}
+	defer a.Close()
+
+	if err := a.SetSystemPrompt("be brief"); err != nil {
+		t.Fatalf("SetSystemPrompt: %v", err)
+	}
+	if err := a.AddSystemMessage("also be accurate"); err != nil {
+		t.Fatalf("AddSystemMessage: %v", err)
+	}
+}