From 512460fdeb506be3a373ee574f0bef068f87418f Mon Sep 17 00:00:00 2001
From: alfred-mk <alfredmwaik@gmail.com>
Date: Fri, 13 Sep 2024 16:02:04 +0300
Subject: [PATCH] Added Custom AtOutput to append CON or END to output

---
 internal/handlers/base.go   | 31 ++++++++++++++++++++++++++-----
 internal/handlers/single.go |  1 +
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/internal/handlers/base.go b/internal/handlers/base.go
index 483aa02..054654b 100644
--- a/internal/handlers/base.go
+++ b/internal/handlers/base.go
@@ -1,13 +1,15 @@
 package handlers
 
 import (
-	"git.defalsify.org/vise.git/engine"
-	"git.defalsify.org/vise.git/resource"
-	"git.defalsify.org/vise.git/persist"
-	"git.defalsify.org/vise.git/db"
+	"io"
+
+	"git.defalsify.org/vise.git/db"
+	"git.defalsify.org/vise.git/engine"
+	"git.defalsify.org/vise.git/persist"
+	"git.defalsify.org/vise.git/resource"
 
-	"git.grassecon.net/urdt/ussd/internal/storage"
 	"git.grassecon.net/urdt/ussd/internal/handlers/ussd"
+	"git.grassecon.net/urdt/ussd/internal/storage"
 )
 
 type BaseSessionHandler struct {
@@ -88,6 +90,25 @@ func(f *BaseSessionHandler) Output(rqs RequestSession) (RequestSession,  error)
 	return rqs, err
 }
 
+func (f *BaseSessionHandler) AtOutput(rqs RequestSession) (RequestSession, error) {
+	var err error
+	var prefix string
+
+	if rqs.Continue {
+		prefix = "CON "
+	} else {
+		prefix = "END "
+	}
+
+	_, err = io.WriteString(rqs.Writer, prefix)
+	if err != nil {
+		return rqs, err
+	}
+
+	_, err = rqs.Engine.WriteResult(rqs.Ctx, rqs.Writer)
+	return rqs, err
+}
+
 func(f *BaseSessionHandler) Reset(rqs RequestSession) (RequestSession, error) {
 	defer f.provider.Put(rqs.Config.SessionId, rqs.Storage)
 	return rqs, rqs.Engine.Finish()
diff --git a/internal/handlers/single.go b/internal/handlers/single.go
index 69a4d1e..a900c14 100644
--- a/internal/handlers/single.go
+++ b/internal/handlers/single.go
@@ -51,6 +51,7 @@ type RequestHandler interface {
 	GetEngine(cfg engine.Config, rs resource.Resource, pe *persist.Persister) engine.Engine 
 	Process(rs RequestSession) (RequestSession, error)
 	Output(rs RequestSession) (RequestSession, error)
+	AtOutput(rs RequestSession) (RequestSession, error)
 	Reset(rs RequestSession) (RequestSession, error)
 	Shutdown()
 }